diff options
| author |  Jianwei Xie <xiejw@google.com> | 2018-01-24 10:02:35 -0800 |
|---|---|---|
| committer |  TensorFlower Gardener <gardener@tensorflow.org> | 2018-01-24 10:06:06 -0800 |
| commit | d9f93c42a50b1f1401d9c186eac0ae8dc9093c3b (patch) | |
| tree | 178d1a692f56580c266139642b5a1d0d155c477e | |
| parent | 7b62a71e2d46c148df7d5704972f4592bc5e0f1b (diff) | |
Merge changes from github.
PiperOrigin-RevId: 183100142
143 files changed, 5265 insertions, 1416 deletions
diff --git a/RELEASE.md b/RELEASE.md index b24e83f053..fdf10407fd 100644 --- a/RELEASE.md +++ b/RELEASE.md @@ -13,46 +13,146 @@ * [TensorFlow Lite](https://github.com/tensorflow/tensorflow/tree/r1.5/tensorflow/contrib/lite) dev preview is now available. * CUDA 9 and cuDNN 7 support. +* Accelerated Linear Algebra (XLA): + * Add `complex64` support to XLA compiler. + * `bfloat` support is now added to XLA infrastructure. + * Make `ClusterSpec` propagation work with XLA devices. + * Use a determinisitic executor to generate XLA graph. +* `tf.contrib`: + * `tf.contrib.distributions`: + * Add `tf.contrib.distributions.Autoregressive`. + * Make `tf.contrib.distributions` QuadratureCompound classes support batch + * Infer `tf.contrib.distributions.RelaxedOneHotCategorical` `dtype` from arguments. + * Make `tf.contrib.distributions` quadrature family parameterized by + `quadrature_grid_and_prob` vs `quadrature_degree`. + * `auto_correlation` added to `tf.contrib.distributions` + * Add `tf.contrib.bayesflow.layers`, a collection of probabilistic (neural) layers. + * Add `tf.contrib.bayesflow.halton_sequence`. + * Add `tf.contrib.data.make_saveable_from_iterator.` + * Add `tf.contrib.data.shuffle_and_repeat`. + * Add new custom transformation: `tf.contrib.data.scan()`. + * `tf.contrib.distributions.bijectors`: + * Add `tf.contrib.distributions.bijectors.MaskedAutoregressiveFlow`. + * Add `tf.contrib.distributions.bijectors.Permute`. + * Add `tf.contrib.distributions.bijectors.Gumbel`. + * Add `tf.contrib.distributions.bijectors.Reshape`. + * Support shape inference (i.e., shapes containing -1) in the Reshape bijector. +* Add `streaming_precision_recall_at_equal_thresholds,` a method for computing + streaming precision and recall with `O(num_thresholds + size of predictions)` + time and space complexity. +* Change `RunConfig` default behavior to not set a random seed, making random + behavior independently random on distributed workers. We expect this to + generally improve training performance. Models that do rely on determinism + should set a random seed explicitly. +* Replaced the implementation of `tf.flags` with `absl.flags`. +* Add support for `CUBLAS_TENSOR_OP_MATH` in fp16 GEMM +* Add support for CUDA on NVIDIA Tegra devices ## Bug Fixes and Other Changes -* `auto_correlation` added to `tf.contrib.distributions`. -* Add `DenseFlipout` probabilistic layer. -* Restandardize `DenseVariational` as simpler template for other probabilistic layers. -* Make `tf.contrib.distributions` QuadratureCompound classes support batch. +* Documentation updates: + * Clarified that you can only install TensorFlow on 64-bit machines. + * Added a short doc explaining how `Estimator`s save checkpoints. + * Add documentation for ops supported by the `tf2xla` bridge. + * Fix minor typos in the doc of `SpaceToDepth` and `DepthToSpace`. + * Updated documentation comments in `mfcc_mel_filterbank.h` and `mfcc.h` to + clarify that the input domain is squared magnitude spectra and the weighting + is done on linear magnitude spectra (sqrt of inputs). + * Change `tf.contrib.distributions` docstring examples to use `tfd` alias + rather than `ds`, `bs`. + * Fix docstring typos in `tf.distributions.bijectors.Bijector`. + * `tf.assert_equal` no longer raises `ValueError.` It now raises + `InvalidArgumentError,` as documented. + * Update Getting Started docs and API intro. +* Google Cloud Storage (GCS): + * Add userspace DNS caching for the GCS client. + * Customize request timeouts for the GCS filesystem. + * Improve GCS filesystem caching. +* Bug Fixes: + * Fix bug where partitioned integer variables got their wrong shapes. Before + * Fix correctness bug in CPU and GPU implementations of Adadelta. + * Fix a bug in `import_meta_graph`'s handling of partitioned variables when + importing into a scope. WARNING: This may break loading checkpoints of + graphs with partitioned variables saved after using `import_meta_graph` with + a non-empty `import_scope` argument. + * Fix bug in offline debugger which prevented viewing events. + * Added the `WorkerService.DeleteWorkerSession` method to the gRPC interface, + to fix a memory leak. Ensure that your master and worker servers are running + the same version of TensorFlow to avoid compatibility issues. + * Fix bug in peephole implementation of BlockLSTM cell. + * Fix bug by casting dtype of `log_det_jacobian` to match `log_prob` in + `TransformedDistribution`. + * Fix a bug in `import_meta_graph`'s handling of partitioned variables when + * Ensure `tf.distributions.Multinomial` doesn't underflow in `log_prob`. + Before this change, all partitions of an integer variable were initialized + with the shape of the unpartitioned variable; after this change they are + initialized correctly. +* Other: + * Add necessary shape util support for bfloat16. + * Add a way to run ops using a step function to MonitoredSession. + * Add `DenseFlipout` probabilistic layer. + * A new flag `ignore_live_threads` is available on train. If set to `True`, it + will ignore threads that remain running when tearing down infrastructure + after successfully completing training, instead of throwing a RuntimeError. + * Restandardize `DenseVariational` as simpler template for other probabilistic + layers. + * `tf.data` now supports `tf.SparseTensor` components in dataset elements. + * It is now possible to iterate over `Tensor`s. + * Allow `SparseSegmentReduction` ops to have missing segment IDs. + * Modify custom export strategy to account for multidimensional sparse float + splits. + * `Conv2D`, `Conv2DBackpropInput`, `Conv2DBackpropFilter` now supports arbitrary + dilations with GPU and cuDNNv6 support. + * `Estimator` now supports `Dataset`: `input_fn` can return a `Dataset` + instead of `Tensor`s. + * Add `RevBlock`, a memory-efficient implementation of reversible residual layers. + * Reduce BFCAllocator internal fragmentation. + * Add `cross_entropy` and `kl_divergence` to `tf.distributions.Distribution`. + * Add `tf.nn.softmax_cross_entropy_with_logits_v2` which enables backprop + w.r.t. the labels. + * GPU back-end now uses `ptxas` to compile generated PTX. + * `BufferAssignment`'s protocol buffer dump is now deterministic. + * Change embedding op to use parallel version of `DynamicStitch`. + * Add support for sparse multidimensional feature columns. + * Speed up the case for sparse float columns that have only 1 value. + * Allow sparse float splits to support multivalent feature columns. + * Add `quantile` to `tf.distributions.TransformedDistribution`. + * Add `NCHW_VECT_C` support for `tf.depth_to_space` on GPU. + * Add `NCHW_VECT_C` support for `tf.space_to_depth` on GPU. + +## API Changes +* Rename `SqueezeDims` attribute to `Axis` in C++ API for Squeeze op. * `Stream::BlockHostUntilDone` now returns Status rather than bool. -* Customize request timeouts for the GCS filesystem. +* Minor refactor: move stats files from `stochastic` to `common` and remove + `stochastic`. ## Thanks to our Contributors This release contains contributions from many people at Google, as well as: -4d55397500, Abdullah Alrasheed, abenmao, Adam Salvail, Aditya Dhulipala, Ag Ramesh, -Akimasa Kimura, Alan Du, Alan Yee, Alexander, Amit Kushwaha, Amy, Andrei Costinescu, -Andrei Nigmatulin, Andrew Erlichson, Andrew Myers, Andrew Stepanov, Androbin, AngryPowman, -Anish Shah, Anton Daitche, Artsiom Chapialiou, asdf2014, Aseem Raj Baranwal, Ash Hall, -Bart Kiers, Batchu Venkat Vishal, ben, Ben Barsdell, Bill Piel, Carl Thomé, Catalin Voss, -Changming Sun, Chengzhi Chen, Chi Zeng, Chris Antaki, Chris Donahue, Chris Oelmueller, -Chris Tava, Clayne Robison, Codrut, Courtial Florian, Dalmo Cirne, Dan J, Darren Garvey, -David Kristoffersson, David Norman, David RöThlisberger, DavidNorman, Dhruv, DimanNe, -Dorokhov, Duncan Mac-Vicar P, EdwardDixon, EMCP, error.d, FAIJUL, Fan Xia, -Francois Xavier, Fred Reiss, Freedom" Koan-Sin Tan, Fritz Obermeyer, Gao, Xiang, -Guenther Schmuelling, Guo Yejun (郭叶军), Hans Gaiser, HectorSVC, Hyungsuk Yoon, -James Pruegsanusak, Jay Young, Jean Wanka, Jeff Carpenter, Jeremy Rutman, Jeroen BéDorf, -Jett Jones, Jimmy Jia, jinghuangintel, jinze1994, JKurland, Joel Hestness, joetoth, -John B Nelson, John Impallomeni, John Lawson, Jonas, Jonathan Dekhtiar, joshkyh, Jun Luan, -Jun Mei, Kai Sasaki, Karl Lessard, karl@kubx.ca, Kb Sriram, Kenichi Ueno, Kevin Slagle, -Kongsea, Lakshay Garg, lhlmgr, Lin Min, liu.guangcong, Loki Der Quaeler, Louie Helm, -lucasmoura, Luke Iwanski, Lyndon White, Mahmoud Abuzaina, Marcel Puyat, Mark Aaron Shirley, -Michele Colombo, MtDersvan, Namrata-Ibm, Nathan Luehr, Naurril, Nayana Thorat, Nicolas Lopez, -Niranjan Hasabnis, Nolan Liu, Nouce, Oliver Hennigh, osdamv, Patrik Erdes, -Patryk Chrabaszcz, Pavel Christof, Penghao Cen, postBG, Qingqing Cao, Qingying Chen, qjivy, -Raphael, Rasmi, raymondxyang, Renze Yu, resec, Roffel, Ruben Vereecken, Ryohei Kuroki, -sandipmgiri, Santiago Castro, Scott Kirkland, Sean Vig, Sebastian Raschka, Sebastian Weiss, -Sergey Kolesnikov, Sergii Khomenko, Shahid, Shivam Kotwalia, Stuart Berg, Sumit Gouthaman, -superzerg, Sven Mayer, tetris, Ti Zhou, Tiago Freitas Pereira, Tian Jin, Tomoaki Oiki, -Vaibhav Sood, vfdev, Vivek Rane, Vladimir Moskva, wangqr, Weber Xie, Will Frey, -Yan Facai (颜发才), yanivbl6, Yaroslav Bulatov, Yixing Lao, Yong Tang, youkaichao, -Yuan (Terry) Tang, Yue Zhang, Yuxin Wu, Ziming Dong, ZxYuan, 黄璞 +Adam Zahran, Ag Ramesh, Alan Lee, Alan Yee, Alex Sergeev, Alexander, Amir H. Jadidinejad, +Amy, Anastasios Doumoulakis, Andrei Costinescu, Andrei Nigmatulin, Anthony Platanios, +Anush Elangovan, arixlin, Armen Donigian, ArtëM Sobolev, Atlas7, Ben Barsdell, Bill Prin, +Bo Wang, Brett Koonce, Cameron Thomas, Carl Thomé, Cem Eteke, cglewis, Changming Sun, +Charles Shenton, Chi-Hung, Chris Donahue, Chris Filo Gorgolewski, Chris Hoyean Song, +Chris Tava, Christian Grail, Christoph Boeddeker, cinqS, Clayne Robison, codrut3, concerttttt, +CQY, Dan Becker, Dan Jarvis, Daniel Zhang, David Norman, dmaclach, Dmitry Trifonov, +Donggeon Lim, dongpilYu, Dr. Kashif Rasul, Edd Wilder-James, Eric Lv, fcharras, Felix Abecassis, +FirefoxMetzger, formath, FredZhang, Gaojin Cao, Gary Deer, Guenther Schmuelling, Hanchen Li, +Hanmin Qin, hannesa2, hyunyoung2, Ilya Edrenkin, Jackson Kontny, Jan, Javier Luraschi, +Jay Young, Jayaram Bobba, Jeff, Jeff Carpenter, Jeremy Sharpe, Jeroen BéDorf, Jimmy Jia, +Jinze Bai, Jiongyan Zhang, Joe Castagneri, Johan Ju, Josh Varty, Julian Niedermeier, +JxKing, Karl Lessard, Kb Sriram, Keven Wang, Koan-Sin Tan, Kyle Mills, lanhin, LevineHuang, +Loki Der Quaeler, Loo Rong Jie, Luke Iwanski, LáSzló Csomor, Mahdi Abavisani, Mahmoud Abuzaina, +ManHyuk, Marek ŠUppa, MathSquared, Mats Linander, Matt Wytock, Matthew Daley, Maximilian Bachl, +mdymczyk, melvyniandrag, Michael Case, Mike Traynor, miqlas, Namrata-Ibm, Nathan Luehr, +Nathan Van Doorn, Noa Ezra, Nolan Liu, Oleg Zabluda, opensourcemattress, Ouwen Huang, +Paul Van Eck, peisong, Peng Yu, PinkySan, pks, powderluv, Qiao Hai-Jun, Qiao Longfei, +Rajendra Arora, Ralph Tang, resec, Robin Richtsfeld, Rohan Varma, Ryohei Kuroki, SaintNazaire, +Samuel He, Sandeep Dcunha, sandipmgiri, Sang Han, scott, Scott Mudge, Se-Won Kim, Simon Perkins, +Simone Cirillo, Steffen Schmitz, Suvojit Manna, Sylvus, Taehoon Lee, Ted Chang, Thomas Deegan, +Till Hoffmann, Tim, Toni Kunic, Toon Verstraelen, Tristan Rice, Urs KöSter, Utkarsh Upadhyay, +Vish (Ishaya) Abrams, Winnie Tsang, Yan Chen, Yan Facai (颜发才), Yi Yang, Yong Tang, +Youssef Hesham, Yuan (Terry) Tang, Zhengsheng Wei, zxcqwe4906, 张志豪, 田传武 We are also grateful to all who filed issues or helped resolve them, asked and answered questions, and were part of inspiring discussions. @@ -60,7 +160,15 @@ answered questions, and were part of inspiring discussions. # Release 1.4.1 ## Bug Fixes and Other Changes -* `LinearClassifier` fix for the Google Cloud Machine Learning Engine. +* `LinearClassifier` fix. + +# Release 1.4.0 + +## Major Features And Improvements +* `tf.keras` is now part of the core TensorFlow API. +* [`tf.data`](http://tensorflow.org/programmers_guide/datasets) is now part of + the core TensorFlow API. + * The API is now subject to backwards compatibility guarantees. # Release 1.4.0 @@ -2,11 +2,11 @@ workspace(name = "org_tensorflow") http_archive( name = "io_bazel_rules_closure", - sha256 = "110fe68753413777944b473c25eed6368c4a0487cee23a7bac1b13cc49d3e257", - strip_prefix = "rules_closure-4af89ef1db659eb41f110df189b67d4cf14073e1", + sha256 = "6691c58a2cd30a86776dd9bb34898b041e37136f2dc7e24cadaeaf599c95c657", + strip_prefix = "rules_closure-08039ba8ca59f64248bb3b6ae016460fe9c9914f", urls = [ - "https://mirror.bazel.build/github.com/bazelbuild/rules_closure/archive/4af89ef1db659eb41f110df189b67d4cf14073e1.tar.gz", - "https://github.com/bazelbuild/rules_closure/archive/4af89ef1db659eb41f110df189b67d4cf14073e1.tar.gz", # 2017-08-28 + "https://mirror.bazel.build/github.com/bazelbuild/rules_closure/archive/08039ba8ca59f64248bb3b6ae016460fe9c9914f.tar.gz", + "https://github.com/bazelbuild/rules_closure/archive/08039ba8ca59f64248bb3b6ae016460fe9c9914f.tar.gz", # 2018-01-16 ], ) diff --git a/tensorflow/BUILD b/tensorflow/BUILD index 1bf7c741b7..2ea0e38c78 100644 --- a/tensorflow/BUILD +++ b/tensorflow/BUILD @@ -662,6 +662,9 @@ filegroup( "//tensorflow/tools/quantization:all_files", "//tensorflow/tools/test:all_files", "//tensorflow/user_ops:all_files", + "//third_party/eigen3:all_files", + "//third_party/fft2d:all_files", + "//third_party/flatbuffers:all_files", "//third_party/hadoop:all_files", "//third_party/sycl:all_files", "//third_party/sycl/sycl:all_files", diff --git a/tensorflow/c/eager/c_api.cc b/tensorflow/c/eager/c_api.cc index e4e33016c2..ecba8c0109 100644 --- a/tensorflow/c/eager/c_api.cc +++ b/tensorflow/c/eager/c_api.cc @@ -190,6 +190,7 @@ TFE_TensorHandle* TFE_TensorHandleCopyToDevice(TFE_TensorHandle* h, bool is_same_device = (srcd == dstd) || (DeviceName(srcd) == DeviceName(dstd)); const bool dst_cpu = IsCPU(dstd); + const bool src_cpu = IsCPU(srcd); if (is_same_device) { return new TFE_TensorHandle(h->t, dst_cpu ? nullptr : dstd); } @@ -213,7 +214,7 @@ TFE_TensorHandle* TFE_TensorHandleCopyToDevice(TFE_TensorHandle* h, return new TFE_TensorHandle(dst, dst_cpu ? nullptr : dstd); } tensorflow::DeviceContext* src_device_context = nullptr; - if (!IsCPU(srcd)) { + if (!src_cpu) { src_device_context = srcd->tensorflow_gpu_device_info()->default_context; } tensorflow::DeviceContext* dst_device_context = nullptr; diff --git a/tensorflow/compiler/xla/tests/dynamic_ops_test.cc b/tensorflow/compiler/xla/tests/dynamic_ops_test.cc index ae3f887240..877dc7db0e 100644 --- a/tensorflow/compiler/xla/tests/dynamic_ops_test.cc +++ b/tensorflow/compiler/xla/tests/dynamic_ops_test.cc @@ -595,6 +595,11 @@ XLA_TEST_F(DynamicUpdateSliceTest, R3ContiguousSingleElement) { // Single element, no wrap. std::vector<int32> operand_shape({4, 5, 2}); RunR3Contiguous<float>(operand_shape, /*index=*/1, /*size=*/1); +} + +XLA_TEST_F(DynamicUpdateSliceTest, R3ContiguousSingleElementBF16) { + // Single element, no wrap. + std::vector<int32> operand_shape({4, 5, 2}); RunR3Contiguous<bfloat16>(operand_shape, /*index=*/1, /*size=*/1); } @@ -602,6 +607,11 @@ XLA_TEST_F(DynamicUpdateSliceTest, R3ContiguousMultipleElements) { // Multiple element, no wrap. std::vector<int32> operand_shape({4, 5, 2}); RunR3Contiguous<float>(operand_shape, /*index=*/1, /*size=*/2); +} + +XLA_TEST_F(DynamicUpdateSliceTest, R3ContiguousMultipleElementsBF16) { + // Multiple element, no wrap. + std::vector<int32> operand_shape({4, 5, 2}); RunR3Contiguous<bfloat16>(operand_shape, /*index=*/1, /*size=*/2); } @@ -609,6 +619,11 @@ XLA_TEST_F(DynamicUpdateSliceTest, R3ContiguousMultipleWrapping) { // Multiple element, wrapping. std::vector<int32> operand_shape({4, 5, 2}); RunR3Contiguous<float>(operand_shape, /*index=*/3, /*size=*/2); +} + +XLA_TEST_F(DynamicUpdateSliceTest, R3ContiguousMultipleWrappingBF16) { + // Multiple element, wrapping. + std::vector<int32> operand_shape({4, 5, 2}); RunR3Contiguous<bfloat16>(operand_shape, /*index=*/3, /*size=*/2); } @@ -616,12 +631,21 @@ XLA_TEST_F(DynamicUpdateSliceTest, R3ContiguousTooLarge) { // Multiple element, update size larger than operand. std::vector<int32> operand_shape({4, 5, 2}); RunR3Contiguous<float>(operand_shape, /*index=*/5, /*size=*/2); +} + +XLA_TEST_F(DynamicUpdateSliceTest, R3ContiguousTooLargeBF16) { + // Multiple element, update size larger than operand. + std::vector<int32> operand_shape({4, 5, 2}); RunR3Contiguous<bfloat16>(operand_shape, /*index=*/5, /*size=*/2); } XLA_TEST_F(DynamicUpdateSliceTest, R3ContiguousUnaligned) { std::vector<int32> operand_shape({3, 123, 247}); RunR3Contiguous<float>(operand_shape, /*index=*/1, /*size=*/1); +} + +XLA_TEST_F(DynamicUpdateSliceTest, R3ContiguousUnalignedBF16) { + std::vector<int32> operand_shape({3, 123, 247}); RunR3Contiguous<bfloat16>(operand_shape, /*index=*/1, /*size=*/1); } @@ -629,6 +653,10 @@ XLA_TEST_F(DynamicUpdateSliceTest, R3ContiguousUnaligned) { XLA_TEST_F(DynamicUpdateSliceTest, DISABLED_ON_GPU(R3ContiguousLarger)) { std::vector<int32> operand_shape({32, 128, 1024}); RunR3Contiguous<float>(operand_shape, /*index=*/7, /*size=*/1); +} + +XLA_TEST_F(DynamicUpdateSliceTest, DISABLED_ON_GPU(R3ContiguousLargerBF16)) { + std::vector<int32> operand_shape({32, 128, 1024}); RunR3Contiguous<bfloat16>(operand_shape, /*index=*/7, /*size=*/1); } diff --git a/tensorflow/contrib/cmake/README.md b/tensorflow/contrib/cmake/README.md index 4be733a280..8f85a75ee4 100644 --- a/tensorflow/contrib/cmake/README.md +++ b/tensorflow/contrib/cmake/README.md @@ -30,7 +30,7 @@ bindings. * CMake version 3.5 or later. -* [Git](http://git-scm.com) +* [Git](https://git-scm.com) * [SWIG](http://www.swig.org/download.html) @@ -48,7 +48,7 @@ bindings. * Microsoft Windows 10 - Microsoft Visual Studio Enterprise 2015 with Visual C++ 2015 - - [Anaconda 4.1.1 (Python 3.5 64-bit)](https://www.continuum.io/downloads) + - [Anaconda 4.1.1 (Python 3.5 64-bit)](https://www.anaconda.com/download/) - [Git for Windows version 2.9.2.windows.1](https://git-scm.com/download/win) - [swigwin-3.0.10](http://www.swig.org/download.html) - [NVidia CUDA Toolkit 8.0](https://developer.nvidia.com/cuda-downloads) diff --git a/tensorflow/contrib/cmake/external/snappy.cmake b/tensorflow/contrib/cmake/external/snappy.cmake index 013b3a862f..fd57734298 100644 --- a/tensorflow/contrib/cmake/external/snappy.cmake +++ b/tensorflow/contrib/cmake/external/snappy.cmake @@ -47,4 +47,4 @@ ExternalProject_Add(snappy ) # actually enables snappy in the source code -add_definitions(-DTF_USE_SNAPPY)
\ No newline at end of file +add_definitions(-DTF_USE_SNAPPY) diff --git a/tensorflow/contrib/cmake/python_modules.txt b/tensorflow/contrib/cmake/python_modules.txt index dec6c513ba..7db454bd83 100644 --- a/tensorflow/contrib/cmake/python_modules.txt +++ b/tensorflow/contrib/cmake/python_modules.txt @@ -1,3 +1,5 @@ +# python_sanity_test.py will complain about invalid or missing entries +# problematic entries can be commented for temporary whitelisting tensorflow tensorflow/core tensorflow/core/example @@ -307,6 +309,8 @@ tensorflow/contrib/metrics tensorflow/contrib/metrics/python tensorflow/contrib/metrics/python/metrics tensorflow/contrib/metrics/python/ops +tensorflow/contrib/mpi_collectives/python +tensorflow/contrib/mpi_collectives/python/ops tensorflow/contrib/model_pruning tensorflow/contrib/model_pruning/examples tensorflow/contrib/model_pruning/examples/cifar10 diff --git a/tensorflow/contrib/cmake/python_sanity_test.py b/tensorflow/contrib/cmake/python_sanity_test.py new file mode 100644 index 0000000000..3be5bd1b23 --- /dev/null +++ b/tensorflow/contrib/cmake/python_sanity_test.py @@ -0,0 +1,124 @@ +# Copyright 2017 The TensorFlow Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +""" +Complain about invalid or missing entries in python_*.txt files. +Problematic entries can be commented for temporary whitelisting. +""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import os +import unittest + + +def abs_path(path): + root = os.path.dirname(__file__) + + for _ in range(3): + root = os.path.join(root, os.pardir) + + path = os.path.join(root, path) + path = os.path.abspath(path) + return path + +def read_entries(test): + with open(abs_path(test.entries_file), "r") as f: + lines = f.readlines() + + lines = [line.strip() for line in lines] + lines = [line for line in lines if line] + + test.entries = [] + test.whitelist = [] + + for line in lines: + # line is comment + if line.startswith('#'): + line = line[1:].strip() + # whitelist entry + if line.startswith('tensorflow/'): + test.whitelist.append(line) + # line has comment -> strip comment + elif line.find('#') != -1: + line = line[:line.find('#')].strip() + test.entries.append(line) + else: + test.entries.append(line) + +def test_invalid_directories(test): + for entry in test.entries: + if not os.path.isdir(abs_path(entry)): + problem = "'" + test.entries_file + "' contains invalid '" + entry + "'" + solution = "Please remove the invalid entry (or add the missing directory)." + raise AssertionError(problem + "\n" + solution) + +def test_missing_directory(test, path): + if path in test.whitelist: + return + + dir_exists = os.path.isdir(abs_path(path)) + entry_exists = path in test.entries + + if dir_exists and not entry_exists: + problem = "'" + test.entries_file + "' is missing '" + path + "'" + solution = "Please add the missing entry (comment to whitelist if needed)." + raise AssertionError(problem + "\n" + solution) + + +class PythonModuleTest(unittest.TestCase): + + def setUp(self): + self.entries_file = "tensorflow/contrib/cmake/python_modules.txt" + read_entries(self) + + def testInvalidEntries(self): + test_invalid_directories(self) + + def testMissingModules(self): + module_names = next(os.walk(abs_path("tensorflow/contrib")))[1] + + for module_name in module_names: + path = "tensorflow/contrib/" + module_name + + test_missing_directory(self, path + "/python") + test_missing_directory(self, path + "/python/ops") + test_missing_directory(self, path + "/python/kernels") + test_missing_directory(self, path + "/python/layers") + + +class PythonProtoTest(unittest.TestCase): + + def setUp(self): + self.entries_file = "tensorflow/contrib/cmake/python_protos.txt" + read_entries(self) + + def testInvalidEntries(self): + test_invalid_directories(self) + + +class PythonProtoCCTest(unittest.TestCase): + + def setUp(self): + self.entries_file = "tensorflow/contrib/cmake/python_protos_cc.txt" + read_entries(self) + + def testInvalidEntries(self): + test_invalid_directories(self) + + +if __name__ == "__main__": + unittest.main() diff --git a/tensorflow/contrib/cmake/tf_core_framework.cmake b/tensorflow/contrib/cmake/tf_core_framework.cmake index 24d7fb82a2..129c208ecd 100644 --- a/tensorflow/contrib/cmake/tf_core_framework.cmake +++ b/tensorflow/contrib/cmake/tf_core_framework.cmake @@ -126,7 +126,9 @@ endfunction() file(GLOB_RECURSE tf_protos_cc_srcs RELATIVE ${tensorflow_source_dir} "${tensorflow_source_dir}/tensorflow/core/*.proto" "${tensorflow_source_dir}/tensorflow/contrib/boosted_trees/proto/*.proto" + "${tensorflow_source_dir}/tensorflow/contrib/tpu/proto/*.proto" ) + RELATIVE_PROTOBUF_GENERATE_CPP(PROTO_SRCS PROTO_HDRS ${tensorflow_source_dir} ${tf_protos_cc_srcs} ) diff --git a/tensorflow/contrib/cmake/tf_python.cmake b/tensorflow/contrib/cmake/tf_python.cmake index b62a031749..8862390d2b 100755 --- a/tensorflow/contrib/cmake/tf_python.cmake +++ b/tensorflow/contrib/cmake/tf_python.cmake @@ -126,7 +126,8 @@ STRING(REGEX REPLACE ";" "\\\\;" python_protos "${python_protos}") STRING(REGEX REPLACE "\n" ";" python_protos "${python_protos}") foreach(python_proto ${python_protos}) - if(NOT python_proto MATCHES "\#") + if(NOT python_proto MATCHES "^\#") + STRING(REGEX REPLACE " *\#.*" "" python_proto "${python_proto}") if(NOT EXISTS "${tensorflow_source_dir}/${python_proto}") message(SEND_ERROR "Python proto directory not found: ${python_proto}") endif() @@ -147,7 +148,8 @@ STRING(REGEX REPLACE ";" "\\\\;" python_protos_cc "${python_protos_cc}") STRING(REGEX REPLACE "\n" ";" python_protos_cc "${python_protos_cc}") foreach(python_proto_cc ${python_protos_cc}) - if(NOT python_proto_cc MATCHES "\#") + if(NOT python_proto_cc MATCHES "^\#") + STRING(REGEX REPLACE " *\#.*" "" python_proto_cc "${python_proto_cc}") if(NOT EXISTS "${tensorflow_source_dir}/${python_proto_cc}") message(SEND_ERROR "Python proto CC directory not found: ${python_proto_cc}") endif() @@ -209,7 +211,8 @@ STRING(REGEX REPLACE ";" "\\\\;" python_modules "${python_modules}") STRING(REGEX REPLACE "\n" ";" python_modules "${python_modules}") foreach(python_module ${python_modules}) - if(NOT python_module MATCHES "\#") + if(NOT python_module MATCHES "^\#") + STRING(REGEX REPLACE " *\#.*" "" python_module "${python_module}") if(NOT EXISTS "${tensorflow_source_dir}/${python_module}") message(SEND_ERROR "Python module not found: ${python_module}") endif() diff --git a/tensorflow/contrib/cudnn_rnn/kernels/cudnn_rnn_ops.cc b/tensorflow/contrib/cudnn_rnn/kernels/cudnn_rnn_ops.cc index 6b0452e7af..ba9686e94e 100644 --- a/tensorflow/contrib/cudnn_rnn/kernels/cudnn_rnn_ops.cc +++ b/tensorflow/contrib/cudnn_rnn/kernels/cudnn_rnn_ops.cc @@ -649,7 +649,7 @@ class CudnnRNNParamsToCanonical<GPUDevice, T> : public CudnnRNNKernelCommon { } const int num_params_per_layer = num_params_ / num_layers / num_dirs; // Number of params applied on inputs. The rest are applied on recurrent - // hiddden states. + // hidden states. const int num_params_input_state = num_params_per_layer / 2; CHECK(num_params_ % (num_layers * num_dirs) == 0) << "Number of params is not a multiple of num_layers * num_dirs."; diff --git a/tensorflow/contrib/cudnn_rnn/python/ops/cudnn_rnn_ops.py b/tensorflow/contrib/cudnn_rnn/python/ops/cudnn_rnn_ops.py index 5f6e3be28f..e87162f0ee 100644 --- a/tensorflow/contrib/cudnn_rnn/python/ops/cudnn_rnn_ops.py +++ b/tensorflow/contrib/cudnn_rnn/python/ops/cudnn_rnn_ops.py @@ -1542,7 +1542,7 @@ class _CudnnRNNNoInputC(_CudnnRNN): params: the parameter buffer created for this model. is_training: whether this operation will be used in training or inference. Returns: - output: the output sequuence. + output: the output sequence. output_h: the final state for h. """ return _cudnn_rnn_no_input_c( diff --git a/tensorflow/contrib/eager/python/g3doc/guide.md b/tensorflow/contrib/eager/python/g3doc/guide.md index 0095ffa0db..7eea93ce1f 100644 --- a/tensorflow/contrib/eager/python/g3doc/guide.md +++ b/tensorflow/contrib/eager/python/g3doc/guide.md @@ -292,7 +292,7 @@ def loss(weight, bias): error = prediction(training_inputs, weight, bias) - training_outputs return tf.reduce_mean(tf.square(error)) -# Function that returns the the derivative of loss with respect to +# Function that returns the derivative of loss with respect to # weight and bias grad = tfe.gradients_function(loss) diff --git a/tensorflow/contrib/estimator/python/estimator/extenders.py b/tensorflow/contrib/estimator/python/estimator/extenders.py index 29c3c73585..c99bf8badb 100644 --- a/tensorflow/contrib/estimator/python/estimator/extenders.py +++ b/tensorflow/contrib/estimator/python/estimator/extenders.py @@ -100,7 +100,7 @@ def add_metrics(estimator, metric_fn): def clip_gradients_by_norm(optimizer, clip_norm): - """Returns an optimizer which clips gradients before appliying them. + """Returns an optimizer which clips gradients before applying them. Example: diff --git a/tensorflow/contrib/ffmpeg/default/ffmpeg_lib.cc b/tensorflow/contrib/ffmpeg/default/ffmpeg_lib.cc index 3c51deefbc..c85b1837ab 100644 --- a/tensorflow/contrib/ffmpeg/default/ffmpeg_lib.cc +++ b/tensorflow/contrib/ffmpeg/default/ffmpeg_lib.cc @@ -82,7 +82,9 @@ std::vector<string> FfmpegVideoCommandLine(const string& input_filename, "-probesize", StrCat(kDefaultProbeSize), "-loglevel", - "error", // Print errors only. + // Info is needed to get the information about stream, etc. + // It is generated to a separate file, not stdout/stderr. + "info", "-hide_banner", // Skip printing build options, version, etc. "-vcodec", "rawvideo", diff --git a/tensorflow/contrib/framework/python/ops/script_ops.py b/tensorflow/contrib/framework/python/ops/script_ops.py index fdbb77bbe4..5d269fefdc 100644 --- a/tensorflow/contrib/framework/python/ops/script_ops.py +++ b/tensorflow/contrib/framework/python/ops/script_ops.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== + """Script Language Operators. See the @{$python/script_ops} guide. @@py_func diff --git a/tensorflow/contrib/gan/python/train.py b/tensorflow/contrib/gan/python/train.py index a32ddd7a06..5d0ac93aec 100644 --- a/tensorflow/contrib/gan/python/train.py +++ b/tensorflow/contrib/gan/python/train.py @@ -279,14 +279,16 @@ def acgan_model( generator_inputs = _convert_tensor_or_l_or_d(generator_inputs) generated_data = generator_fn(generator_inputs) with variable_scope.variable_scope(discriminator_scope) as dis_scope: - (discriminator_gen_outputs, discriminator_gen_classification_logits - ) = _validate_acgan_discriminator_outputs( - discriminator_fn(generated_data, generator_inputs)) + with ops.name_scope(dis_scope.name+'/generated/'): + (discriminator_gen_outputs, discriminator_gen_classification_logits + ) = _validate_acgan_discriminator_outputs( + discriminator_fn(generated_data, generator_inputs)) with variable_scope.variable_scope(dis_scope, reuse=True): - real_data = ops.convert_to_tensor(real_data) - (discriminator_real_outputs, discriminator_real_classification_logits - ) = _validate_acgan_discriminator_outputs( - discriminator_fn(real_data, generator_inputs)) + with ops.name_scope(dis_scope.name+'/real/'): + real_data = ops.convert_to_tensor(real_data) + (discriminator_real_outputs, discriminator_real_classification_logits + ) = _validate_acgan_discriminator_outputs( + discriminator_fn(real_data, generator_inputs)) if check_shapes: if not generated_data.shape.is_compatible_with(real_data.shape): raise ValueError( diff --git a/tensorflow/contrib/layers/python/layers/layers.py b/tensorflow/contrib/layers/python/layers/layers.py index f3229a1605..ef2b673074 100644 --- a/tensorflow/contrib/layers/python/layers/layers.py +++ b/tensorflow/contrib/layers/python/layers/layers.py @@ -479,8 +479,12 @@ def batch_norm(inputs, Sergey Ioffe, Christian Szegedy - Can be used as a normalizer function for conv2d and fully_connected. - + Can be used as a normalizer function for conv2d and fully_connected. The + normalization is over all but the last dimension if `data_format` is `NHWC` + and all but the second dimension if `data_format` is `NCHW`. In case of a 2D + tensor this corresponds to the batch dimension, while in case of a 4D tensor this + corresponds to the batch and space dimensions. + Note: when training, the moving_mean and moving_variance need to be updated. By default the update ops are placed in `tf.GraphKeys.UPDATE_OPS`, so they need to be added as a dependency to the `train_op`. For example: diff --git a/tensorflow/contrib/lite/download_dependencies.sh b/tensorflow/contrib/lite/download_dependencies.sh index 362e5bee25..e1b7b3613a 100755 --- a/tensorflow/contrib/lite/download_dependencies.sh +++ b/tensorflow/contrib/lite/download_dependencies.sh @@ -22,7 +22,14 @@ cd "$SCRIPT_DIR/../../.." DOWNLOADS_DIR=tensorflow/contrib/lite/downloads BZL_FILE_PATH=tensorflow/workspace.bzl -EIGEN_URL="$(grep -o 'http.*bitbucket.org/eigen/eigen/get/.*tar\.gz' "${BZL_FILE_PATH}" | grep -v bazel-mirror | head -n1)" +# Ensure it is being run from repo root +if [ ! -f $BZL_FILE_PATH ]; then + echo "Could not find ${BZL_FILE_PATH}": + echo "Likely you are not running this from the root directory of the repository."; + exit 1; +fi + +EIGEN_URL="$(grep -o 'http.*bitbucket.org/eigen/eigen/get/.*tar\.gz' "${BZL_FILE_PATH}" | grep -v mirror.bazel | head -n1)" GEMMLOWP_URL="$(grep -o 'https://mirror.bazel.build/github.com/google/gemmlowp/.*zip' "${BZL_FILE_PATH}" | head -n1)" GOOGLETEST_URL="https://github.com/google/googletest/archive/release-1.8.0.tar.gz" ABSL_URL="$(grep -o 'https://github.com/abseil/abseil-cpp/.*tar.gz' "${BZL_FILE_PATH}" | head -n1)" diff --git a/tensorflow/contrib/lite/kernels/gather.cc b/tensorflow/contrib/lite/kernels/gather.cc index f8df797daf..0e4187d1ea 100644 --- a/tensorflow/contrib/lite/kernels/gather.cc +++ b/tensorflow/contrib/lite/kernels/gather.cc @@ -42,9 +42,10 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) { TF_LITE_ENSURE_EQ(context, positions->type, kTfLiteInt32); // Check that input and output types match. TF_LITE_ENSURE_EQ(context, input->type, output->type); - // TODO(mgubin): only 1D positions are currently supported. - TF_LITE_ENSURE_EQ(context, NumDimensions(positions), 1); + // TODO(mgubin): only 0D or 1D positions are currently supported. + TF_LITE_ENSURE(context, NumDimensions(positions) <= 1); // TODO(mgubin): Only default axis == 0 is supported. + TF_LITE_ENSURE_EQ(context, params->axis, 0); // Check conditions for different types. switch (input->type) { case kTfLiteFloat32: @@ -64,7 +65,7 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) { } const int num_dimensions = NumDimensions(input) + NumDimensions(positions) - 1; - TF_LITE_ENSURE(context, params->axis < num_dimensions); + TF_LITE_ENSURE(context, params->axis <= num_dimensions); TfLiteIntArray* output_shape = TfLiteIntArrayCreate(num_dimensions); int output_index = 0; for (int i = 0; i < params->axis; ++i) { diff --git a/tensorflow/contrib/lite/kernels/gather_test.cc b/tensorflow/contrib/lite/kernels/gather_test.cc index 6343d3b4ef..658d977b8d 100644 --- a/tensorflow/contrib/lite/kernels/gather_test.cc +++ b/tensorflow/contrib/lite/kernels/gather_test.cc @@ -48,8 +48,8 @@ class GatherOpModel : public SingleOpModel { PopulateStringTensor(input_, data); } - void SetPositions(std::initializer_list<int32> data) { - PopulateTensor<int32>(positions_, data); + void SetPositions(std::initializer_list<int> data) { + PopulateTensor<int>(positions_, data); } std::vector<float> GetOutputFloat() { return ExtractVector<float>(output_); } @@ -76,6 +76,29 @@ TEST(GatherOpTest, Shuffle) { ElementsAreArray(ArrayFloatNear({0.7, 0.8, -2, 0.2}))); } +TEST(GatherOpTest, Test0DIndex) { + GatherOpModel m({2, 2}, TensorType_FLOAT32, {}); + m.SetInputFloat({-2.0, 0.2, 0.7, 0.8}); + m.SetPositions({1}); + m.Invoke(); + EXPECT_THAT(m.GetOutputFloat(), + ElementsAreArray(ArrayFloatNear({0.7, 0.8}))); + EXPECT_THAT(m.GetOutputShape(), + ElementsAreArray({2})); +} + +TEST(GatherOpTest, Test0DIndexWith0DResult) { + // 0D tensor is special case in current TFLite. Test it once to make sure + // existing workarounds are fine with it. + GatherOpModel m({3}, TensorType_FLOAT32, {}); + m.SetInputFloat({1.0, 2.0, 3.0}); + m.SetPositions({1}); + m.Invoke(); + EXPECT_THAT(m.GetOutputFloat(), + ElementsAreArray(ArrayFloatNear({2.0}))); + EXPECT_TRUE(m.GetOutputShape().empty()); +} + TEST(FloatGatherOpTest, Duplicate) { GatherOpModel m({1, 2, 2}, TensorType_FLOAT32, {2}); m.SetInputFloat({-2.0, 0.2, 0.7, 0.8}); diff --git a/tensorflow/contrib/lite/models/testdata/g3doc/README.md b/tensorflow/contrib/lite/models/testdata/g3doc/README.md index 667a588383..1c47e00aae 100644 --- a/tensorflow/contrib/lite/models/testdata/g3doc/README.md +++ b/tensorflow/contrib/lite/models/testdata/g3doc/README.md @@ -53,7 +53,7 @@ with the corresponding parameters as shown in the figure. ### Automatic Speech Recognizer (ASR) Acoustic Model (AM) The acoustic model for automatic speech recognition is the neural network model -for matching phonemes to the input autio features. It generates posterior +for matching phonemes to the input audio features. It generates posterior probabilities of phonemes from speech frontend features (log-mel filterbanks). It has an input size of 320 (float), an output size of 42 (float), five LSTM layers and one fully connected layers with a Softmax activation function, with @@ -68,7 +68,7 @@ for predicting the probability of a word given previous words in a sentence. It generates posterior probabilities of the next word based from a sequence of words. The words are encoded as indices in a fixed size dictionary. The model has two inputs both of size one (integer): the current word index and -next word index, an output size of one (float): the log probability. It consits +next word index, an output size of one (float): the log probability. It consists of three embedding layer, three LSTM layers, followed by a multiplication, a fully connected layers and an addition. The corresponding parameters as shown in the figure. diff --git a/tensorflow/contrib/lite/nnapi/NeuralNetworksShim.h b/tensorflow/contrib/lite/nnapi/NeuralNetworksShim.h index 3cda4bcccc..7019c29959 100644 --- a/tensorflow/contrib/lite/nnapi/NeuralNetworksShim.h +++ b/tensorflow/contrib/lite/nnapi/NeuralNetworksShim.h @@ -370,7 +370,7 @@ enum { * Looks up items from a given tensor. * * Each item in the output is a raw copy of the corresponding item in - * the input “values”. If the the given “lookup” indices are out of bounds, + * the input “values”. If the given “lookup” indices are out of bounds, * the op will fail and an error will be reported. * * Inputs: diff --git a/tensorflow/contrib/lite/testing/generate_examples.py b/tensorflow/contrib/lite/testing/generate_examples.py index 6204471e52..c225cd4f00 100644 --- a/tensorflow/contrib/lite/testing/generate_examples.py +++ b/tensorflow/contrib/lite/testing/generate_examples.py @@ -1170,7 +1170,7 @@ def make_pad_tests(zip_path): def make_reshape_tests(zip_path): """Make a set of tests to do reshape.""" - # Alll shapes below are suitable for tensors with 420 elements. + # All shapes below are suitable for tensors with 420 elements. test_parameters = [{ "dtype": [tf.float32, tf.int32], "input_shape": [[3, 4, 5, 7], [4, 105], [21, 5, 2, 2], [420]], diff --git a/tensorflow/contrib/lite/toco/g3doc/cmdline_reference.md b/tensorflow/contrib/lite/toco/g3doc/cmdline_reference.md index 4776741ab9..5e07795223 100644 --- a/tensorflow/contrib/lite/toco/g3doc/cmdline_reference.md +++ b/tensorflow/contrib/lite/toco/g3doc/cmdline_reference.md @@ -229,7 +229,7 @@ additional information about the multiple input arrays: well-formed quantized representation of these graphs. Such graphs should be fixed, but as a temporary work-around, setting this reorder_across_fake_quant flag allows the converter to perform necessary - graph transformaitons on them, at the cost of no longer faithfully matching + graph transformations on them, at the cost of no longer faithfully matching inference and training arithmetic. ### Logging flags diff --git a/tensorflow/contrib/lite/tools/BUILD b/tensorflow/contrib/lite/tools/BUILD index 6e2d633765..20df905270 100644 --- a/tensorflow/contrib/lite/tools/BUILD +++ b/tensorflow/contrib/lite/tools/BUILD @@ -27,6 +27,27 @@ tf_cc_binary( ], ) +tf_cc_binary( + name = "benchmark_model", + srcs = ["benchmark_model.cc"], + linkopts = select({ + "//tensorflow:android": [ + "-pie", + "-landroid", + "-lm", + "-z defs", + "-Wl,--exclude-libs,ALL", # Exclude syms in all libs from auto export + ], + "//conditions:default": [], + }), + deps = [ + ":mutable_op_resolver", + "//tensorflow/contrib/lite:framework", + "//tensorflow/contrib/lite:string_util", + "//tensorflow/contrib/lite/kernels:builtin_ops", + ], +) + cc_library( name = "gen_op_registration", srcs = ["gen_op_registration.cc"], diff --git a/tensorflow/contrib/makefile/download_dependencies.sh b/tensorflow/contrib/makefile/download_dependencies.sh index 0a47f50c43..4ae18b2cef 100755 --- a/tensorflow/contrib/makefile/download_dependencies.sh +++ b/tensorflow/contrib/makefile/download_dependencies.sh @@ -63,12 +63,17 @@ download_and_extract() { elif [[ "${url}" == *zip ]]; then tempdir=$(mktemp -d) tempdir2=$(mktemp -d) - wget -P ${tempdir} ${url} - unzip ${tempdir}/* -d ${tempdir2} + if [[ "$OSTYPE" == "darwin"* ]]; then + # macOS (AKA darwin) doesn't have wget. + (cd "${tempdir}"; curl --remote-name --silent --location "${url}") + else + wget -P "${tempdir}" "${url}" + fi + unzip "${tempdir}"/* -d "${tempdir2}" # unzip has no strip components, so unzip to a temp dir, and move the files # we want from the tempdir to destination. - cp -R ${tempdir2}/*/* ${dir}/ - rm -rf ${tempdir2} ${tempdir} + cp -R "${tempdir2}"/*/* "${dir}"/ + rm -rf "${tempdir2}" "${tempdir}" fi # Delete any potential BUILD files, which would interfere with Bazel builds. diff --git a/tensorflow/contrib/opt/BUILD b/tensorflow/contrib/opt/BUILD index 9c961f2b9c..827279bd47 100644 --- a/tensorflow/contrib/opt/BUILD +++ b/tensorflow/contrib/opt/BUILD @@ -19,6 +19,7 @@ py_library( "python/training/elastic_average_optimizer.py", "python/training/external_optimizer.py", "python/training/lazy_adam_optimizer.py", + "python/training/model_average_optimizer.py", "python/training/moving_average_optimizer.py", "python/training/multitask_optimizer_wrapper.py", "python/training/nadam_optimizer.py", @@ -193,6 +194,27 @@ tf_py_test( ], ) +tf_py_test( + name = "model_average_optimizer_test", + srcs = ["python/training/model_average_optimizer_test.py"], + additional_deps = [ + ":opt_py", + "//tensorflow/python:client", + "//tensorflow/python:client_testlib", + "//tensorflow/python:array_ops", + "//tensorflow/python:variables", + "//tensorflow/python:framework", + "//tensorflow/python:platform", + "//tensorflow/python:training", + "//tensorflow/python:ops", + "//tensorflow/python:framework_for_generated_wrappers", + "//third_party/py/numpy", + ], + tags = [ + "notap", # This test launches local server. + ], +) + py_test( name = "sign_decay_test", srcs = ["python/training/sign_decay_test.py"], diff --git a/tensorflow/contrib/opt/__init__.py b/tensorflow/contrib/opt/__init__.py index 90d2f92462..6c1bb1adc0 100644 --- a/tensorflow/contrib/opt/__init__.py +++ b/tensorflow/contrib/opt/__init__.py @@ -29,6 +29,7 @@ from tensorflow.contrib.opt.python.training.nadam_optimizer import * from tensorflow.contrib.opt.python.training.powersign import * from tensorflow.contrib.opt.python.training.variable_clipping_optimizer import * from tensorflow.contrib.opt.python.training.elastic_average_optimizer import * +from tensorflow.contrib.opt.python.training.model_average_optimizer import * # pylint: enable=wildcard-import from tensorflow.python.util.all_util import remove_undocumented @@ -48,7 +49,9 @@ _allowed_symbols = [ 'MultitaskOptimizerWrapper', 'clip_gradients_by_global_norm', 'ElasticAverageOptimizer', - 'ElasticAverageCustomGetter' + 'ElasticAverageCustomGetter', + 'ModelAverageOptimizer', + 'ModelAverageCustomGetter' ] remove_undocumented(__name__, _allowed_symbols) diff --git a/tensorflow/contrib/opt/python/training/model_average_optimizer.py b/tensorflow/contrib/opt/python/training/model_average_optimizer.py new file mode 100644 index 0000000000..47509ecca6 --- /dev/null +++ b/tensorflow/contrib/opt/python/training/model_average_optimizer.py @@ -0,0 +1,299 @@ +# 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. +# ============================================================================== + +"""Wrapper optimizer for Model Average """ +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +from tensorflow.python.framework import ops +from tensorflow.python.framework import dtypes +from tensorflow.python.framework import constant_op +from tensorflow.python.training import optimizer +from tensorflow.python.training import session_run_hook +from tensorflow.python.ops import math_ops +from tensorflow.python.ops import control_flow_ops +from tensorflow.python.ops import variable_scope +from tensorflow.python.ops import variables +from tensorflow.python.ops import state_ops +from tensorflow.python.ops import array_ops +from tensorflow.python.ops import data_flow_ops + +GLOBAL_VARIABLE_NAME = 'global_center_variable' + + +class ModelAverageCustomGetter(object): + """Custom_getter class is used to do: + 1. Change trainable variables to local collection and place them at worker + device + 2. Generate global variables + Notice that the class should be used with tf.replica_device_setter, + so that the global center variables and global step variable can be placed + at ps device. Besides, use 'tf.get_variable' instead of 'tf.Variable' to + use this custom getter. + + For example, + ma_custom_getter = ModelAverageCustomGetter(worker_device) + with tf.device( + tf.train.replica_device_setter( + worker_device=worker_device, + ps_device="/job:ps/cpu:0", + cluster=cluster)), + tf.variable_scope('',custom_getter=ma_custom_getter): + hid_w = tf.get_variable( + initializer=tf.truncated_normal( + [IMAGE_PIXELS * IMAGE_PIXELS, FLAGS.hidden_units], + stddev=1.0 / IMAGE_PIXELS), + name="hid_w") + hid_b = tf.get_variable(initializer=tf.zeros([FLAGS.hidden_units]), + name="hid_b") + """ + + def __init__(self, worker_device): + """Create a new `ElasticAverageCustomGetter`. + + Args: + worker_device: String. Name of the `worker` job. + """ + self._worker_device = worker_device + self._local_2_global = {} + + def __call__(self, getter, name, trainable, collections, *args, **kwargs): + if trainable: + with ops.device(self._worker_device): + local_var = getter(name, trainable=True, + collections=[ops.GraphKeys.LOCAL_VARIABLES], + *args, **kwargs) + + global_variable = variable_scope.variable( + name='%s/%s' % (GLOBAL_VARIABLE_NAME, name), + initial_value=local_var.initialized_value(), + trainable=False, + collections=[ops.GraphKeys.GLOBAL_VARIABLES]) + + self._local_2_global[local_var] = global_variable + return local_var + else: + return getter(name, trainable, collections, *args, **kwargs) + + +class ModelAverageOptimizer(optimizer.Optimizer): + """Wrapper optimizer that implements the Model Average algorithm. + This is a sync optimizer. During the training, each worker will update + the local variables and maintains its own local_step, which starts from 0 + and is incremented by 1 after each update of local variables. Whenever the + interval_steps divides the local step, the local variables from all the + workers will be averaged and assigned to global center variables. Then the + local variables will be assigned by global center variables. + """ + + def __init__( + self, + opt, + num_worker, + is_chief, + ma_custom_getter, + interval_steps=100, + use_locking=True, + name="ModelAverageOptimizer"): + """Construct a new model average optimizer. + + Args: + opt: The actual optimizer that will be used to update local variables + num_worker: The number of workers + is_chief: whether chief worker + ma_custom_getter: ModelAverageCustomGetter + interval_steps: An int point value to controls the frequency of the + average of local variables + use_locking: If True use locks for update operations + name: string. Optional name of the returned operation + """ + super(ModelAverageOptimizer, self).__init__(use_locking, name) + self._opt = opt + self._num_worker = num_worker + self._is_chief = is_chief + self._local_2_global = ma_custom_getter._local_2_global + self._interval_steps = interval_steps + self._accumulator_list = [] + self._chief_init_op = None + + self._local_step = variable_scope.get_variable( + initializer=0, + trainable=False, + collections=[ops.GraphKeys.LOCAL_VARIABLES], + name="local_step") + + self._opt._prepare() + + def compute_gradients(self, *args, **kwargs): + """Compute gradients of "loss" for the variables in "var_list". + + This simply wraps the compute_gradients() from the real optimizer. + + Args: + *args: Arguments for compute_gradients(). + **kwargs: Keyword arguments for compute_gradients(). + + Returns: + A list of (gradient, variable) pairs. + """ + return self._opt.compute_gradients(*args, **kwargs) + + def _local_vars_update(self, var_list): + """Get the update ops for the local variables in "var_list". + + Args: + var_list: Optional list or tuple of 'tf.Variable' to update + + Returns: + An update op + """ + if not var_list: + raise ValueError( + 'The list of local_variables should not be empty') + update_ops = [] + global_center_vars = [self._local_2_global[var] for var in var_list] + for lvar, gvar in zip(var_list, global_center_vars): + with ops.device(lvar.device): + update_ops.append(state_ops.assign(lvar, gvar.read_value())) + return control_flow_ops.group(*(update_ops)) + + def apply_gradients(self, grads_and_vars, global_step=None, name=None): + """Apply gradients to variables. + + This contains most of the synchronization implementation and also wraps the + apply_gradients() from the real optimizer. The chief work updates global + variables. + + Args: + grads_and_vars: List of (gradient, variable) pairs as returned by + compute_gradients(). + global_step: Optional Variable to increment by one after the + variables have been updated. + name: Optional name for the returned operation. Default to the + name passed to the Optimizer constructor. + + Returns: + A conditional 'Operation' that update both local and global variables or + just local variables + + Raises: + ValueError: If the grads_and_vars is empty. + ValueError: If global step is not provided, the staleness cannot be + checked. + """ + + # update local variables + if not grads_and_vars: + raise ValueError("Must supply at least one variable") + if global_step is None: + raise ValueError("Global step is required") + + apply_updates = self._opt.apply_gradients(grads_and_vars) + with ops.control_dependencies([apply_updates]): + local_update = state_ops.assign_add( + self._local_step, 1, name='local_step_update').op + + # update global variables. + def _Update_global_variables(): + local_vars = [v for g, v in grads_and_vars if g is not None] + global_vars = [self._local_2_global[v] for v in local_vars] + # sync queue + with ops.colocate_with(global_step): + sync_queue = data_flow_ops.FIFOQueue(-1, [dtypes.bool], shapes=[[]], + shared_name='sync_queue') + train_ops = [] + aggregated_vars = [] + with ops.name_scope(None, self._name + '/global'): + for var, gvar in zip(local_vars, global_vars): + with ops.device(gvar.device): + if isinstance(var._ref(), ops.Tensor): + var_accum = data_flow_ops.ConditionalAccumulator( + var.dtype, + shape=var.get_shape(), + shared_name=gvar.name + "/var_accum") + train_ops.append( + var_accum.apply_grad(var._ref(), local_step=global_step)) + aggregated_vars.append(var_accum.take_grad(self._num_worker)) + else: + raise ValueError("Unknown local variable type!") + self._accumulator_list.append((var_accum, gvar.device)) + # chief worker updates global vars and enqueues tokens to the sync queue + if self._is_chief: + update_ops = [] + with ops.control_dependencies(train_ops): + for avg_var, gvar in zip(aggregated_vars, global_vars): + with ops.device(gvar.device): + update_ops.append(state_ops.assign(gvar, avg_var)) + with ops.device(global_step.device): + update_ops.append(state_ops.assign_add(global_step, 1)) + with ops.control_dependencies(update_ops), ops.device( + global_step.device): + tokens = array_ops.fill([self._num_worker - 1], + constant_op.constant(False)) + sync_op = sync_queue.enqueue_many(tokens) + else: + with ops.control_dependencies(train_ops), ops.device( + global_step.device): + sync_op = sync_queue.dequeue() + + with ops.control_dependencies([sync_op]): + local_update_op = self._local_vars_update(local_vars) + return local_update_op + + with ops.control_dependencies([local_update]): + condition = math_ops.equal(math_ops.mod( + self._local_step, self._interval_steps), 0) + conditional_update = control_flow_ops.cond( + condition, _Update_global_variables, control_flow_ops.no_op) + + chief_init_ops = [] + for accum, dev in self._accumulator_list: + with ops.device(dev): + chief_init_ops.append( + accum.set_global_step( + global_step, name="SetGlobalStep")) + self._chief_init_op = control_flow_ops.group(*(chief_init_ops)) + + return conditional_update + + def get_init_op(self): + """Returns the op to let all the local variables equal to the global + variables before the training begins""" + return self._local_vars_update(variables.trainable_variables()) + + def make_session_run_hook(self): + """Creates a hook to handle ModelAverage ops such as initialization.""" + return _ModelAverageOptimizerHook(self, self._is_chief) + + +class _ModelAverageOptimizerHook(session_run_hook.SessionRunHook): + def __init__(self, ma_optimizer, is_chief): + """Creates hook to handle ModelAverageOptimizer initialization ops. + + Args: + ea_optimizer: `ModelAverageOptimizer` which this hook will initialize. + is_chief: `Bool`, whether is this a chief replica or not. + """ + self._ma_optimizer = ma_optimizer + self._is_chief = is_chief + + def begin(self): + self._local_init_op = variables.local_variables_initializer() + self._global_init_op = None + if self._is_chief: + self._global_init_op = variables.global_variables_initializer() + self._chief_init_op = self._ma_optimizer._chief_init_op + self._variable_init_op = self._ma_optimizer.get_init_op() diff --git a/tensorflow/contrib/opt/python/training/model_average_optimizer_test.py b/tensorflow/contrib/opt/python/training/model_average_optimizer_test.py new file mode 100644 index 0000000000..a73aa772bb --- /dev/null +++ b/tensorflow/contrib/opt/python/training/model_average_optimizer_test.py @@ -0,0 +1,200 @@ +# Copyright 2017 The TensorFlow Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT 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 ModelAverageOptimizer.""" +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import portpicker +from tensorflow.python.framework import constant_op +from tensorflow.python.framework import ops +from tensorflow.python.ops import variables +from tensorflow.python.platform import test +from tensorflow.python.training import gradient_descent +from tensorflow.python.training import server_lib +from tensorflow.python.training import training +from tensorflow.python.training import training_util +from tensorflow.python.ops import variable_scope +from tensorflow.python.training import device_setter +from tensorflow.contrib.opt.python.training.model_average_optimizer import \ + ModelAverageOptimizer, ModelAverageCustomGetter, GLOBAL_VARIABLE_NAME + + +def create_local_cluster(num_workers, num_ps, protocol="grpc"): + """Create local GRPC servers and return them.""" + worker_ports = [portpicker.pick_unused_port() for _ in range(num_workers)] + ps_ports = [portpicker.pick_unused_port() for _ in range(num_ps)] + cluster_dict = { + "worker": ["localhost:%s" % port for port in worker_ports], + "ps": ["localhost:%s" % port for port in ps_ports] + } + cs = server_lib.ClusterSpec(cluster_dict) + + workers = [ + server_lib.Server( + cs, job_name="worker", protocol=protocol, task_index=ix, start=True) + for ix in range(num_workers) + ] + ps_servers = [ + server_lib.Server( + cs, job_name="ps", protocol=protocol, task_index=ix, start=True) + for ix in range(num_ps) + ] + + return cluster_dict, workers, ps_servers + + +# Creates the workers and return their sessions, graphs, train_ops. +# Cheif worker will update at last +def _get_workers(num_workers, steps, workers): + sessions = [] + graphs = [] + train_ops = [] + for worker_id in range(num_workers): + graph = ops.Graph() + is_chief = (worker_id == 0) + with graph.as_default(): + worker_device = "/job:worker/task:%d/cpu:0" % (worker_id) + ma_coustom = ModelAverageCustomGetter( + worker_device=worker_device) + with variable_scope.variable_scope('', + custom_getter=ma_coustom), ops.device( + device_setter.replica_device_setter(worker_device=worker_device, + ps_device="/job:ps/task:0/cpu:0", + ps_tasks=1)): + + global_step = variables.Variable(0, name='global_step', + trainable=False) + var_0 = variable_scope.get_variable(initializer=0.0, name="v0") + var_1 = variable_scope.get_variable(initializer=1.0, name="v1") + + with ops.device("/job:worker/task:" + str(worker_id)): + if worker_id == 0: + grads_0 = constant_op.constant(-1.0) + grads_1 = constant_op.constant(-1.0) + else: + grads_0 = constant_op.constant(-2.0) + grads_1 = constant_op.constant(-2.0) + sgd_opt = gradient_descent.GradientDescentOptimizer(1.0) + opt = ModelAverageOptimizer( + opt=sgd_opt, + num_worker=num_workers, + ma_custom_getter=ma_coustom, + is_chief=is_chief, + interval_steps=steps + ) + train_op = [ + opt.apply_gradients( + [[grads_0, var_0], + [grads_1, var_1]], global_step) + ] + easgd_hook = opt.make_session_run_hook() + # Creates MonitoredSession + sess = training.MonitoredTrainingSession(workers[worker_id].target, + hooks=[easgd_hook]) + + sessions.append(sess) + graphs.append(graph) + train_ops.append(train_op) + return sessions, graphs, train_ops + + +class ModelAverageOptimizerTest(test.TestCase): + def _run(self, train_op, sess): + sess.run(train_op) + + def test1Workers2Period(self): + num_workers = 2 + steps = 2 + num_ps = 1 + cluster, workers, _ = create_local_cluster(num_workers=num_workers, + num_ps=num_ps) + + sessions, graphs, train_ops = _get_workers(num_workers, + steps, + workers) + + var_0 = graphs[0].get_tensor_by_name('v0:0') + var_1 = graphs[0].get_tensor_by_name('v1:0') + global_step = training_util.get_global_step(graphs[0]) + global_var_0 = graphs[0].get_tensor_by_name(GLOBAL_VARIABLE_NAME + "/v0:0") + global_var_1 = graphs[0].get_tensor_by_name(GLOBAL_VARIABLE_NAME + "/v1:0") + + # Verify the initialized value. + self.assertAllEqual(0.0, sessions[0].run(var_0)) + self.assertAllEqual(1.0, sessions[0].run(var_1)) + self.assertAllEqual(0.0, sessions[0].run(global_var_0)) + self.assertAllEqual(1.0, sessions[0].run(global_var_1)) + self.assertAllEqual(0, sessions[0].run(global_step)) + + sessions[0].run(train_ops[0]) + sessions[1].run(train_ops[1]) + + self.assertAllEqual(1.0, sessions[0].run(var_0)) + self.assertAllEqual(2.0, sessions[0].run(var_1)) + self.assertAllEqual(0.0, sessions[0].run(global_var_0)) + self.assertAllEqual(1.0, sessions[0].run(global_var_1)) + self.assertAllEqual(0, sessions[0].run(global_step)) + + # iteration 2, global varibale update + thread_0 = self.checkedThread( + target=self._run, args=(train_ops[0], sessions[0])) + thread_1 = self.checkedThread( + target=self._run, args=(train_ops[1], sessions[1])) + thread_0.start() + thread_1.start() + thread_0.join() + thread_1.join() + + self.assertAllEqual(3.0, sessions[0].run(var_0)) + self.assertAllEqual(4.0, sessions[0].run(var_1)) + self.assertAllEqual(3.0, sessions[0].run(global_var_0)) + self.assertAllEqual(4.0, sessions[0].run(global_var_1)) + self.assertAllEqual(1, sessions[0].run(global_step)) + + # iteration 3 + sessions[0].run(train_ops[0]) + + self.assertAllEqual(4.0, sessions[0].run(var_0)) + self.assertAllEqual(5.0, sessions[0].run(var_1)) + self.assertAllEqual(3.0, sessions[0].run(global_var_0)) + self.assertAllEqual(4.0, sessions[0].run(global_var_1)) + self.assertAllEqual(1, sessions[0].run(global_step)) + + def testPS2TasksWithClusterSpecClass(self): + cluster_spec = server_lib.ClusterSpec({ + "ps": ["ps0:2222", "ps1:2222"], + "worker": ["worker0:2222", "worker1:2222", "worker2:2222"] + }) + worker_device = "/job:worker/task:0" + ma_coustom = ModelAverageCustomGetter( + worker_device=worker_device) + from tensorflow.python.training import device_setter + with ops.device( + device_setter.replica_device_setter(cluster=cluster_spec, + worker_device=worker_device, + ps_device="/job:ps")), \ + variable_scope.variable_scope('', custom_getter=ma_coustom): + v = variable_scope.get_variable(initializer=[1, 2], name="v") + w = variable_scope.get_variable(initializer=[2, 1], name='w') + v_g, w_g = ma_coustom._local_2_global[v], ma_coustom._local_2_global[w] + self.assertDeviceEqual("/job:worker/task:0", v.device) + self.assertDeviceEqual("job:ps/task:0", v_g.device) + self.assertDeviceEqual("/job:worker/task:0", w.device) + self.assertDeviceEqual("job:ps/task:1", w_g.device) + + +if __name__ == '__main__': + test.main() diff --git a/tensorflow/contrib/periodic_resample/BUILD b/tensorflow/contrib/periodic_resample/BUILD index 71582f9c9a..bd9078ae76 100644 --- a/tensorflow/contrib/periodic_resample/BUILD +++ b/tensorflow/contrib/periodic_resample/BUILD @@ -6,6 +6,7 @@ exports_files(["LICENSE"]) load( "//tensorflow:tensorflow.bzl", + "py_test", "tf_gen_op_libs", "tf_custom_op_library", "tf_custom_op_py_library", @@ -64,11 +65,28 @@ py_library( "python/__init__.py", ], srcs_version = "PY2AND3", + tags = [ + "notap", + ], deps = [ ":periodic_resample_op_py", ], ) +py_test( + name = "periodic_resample_op_test", + srcs = ["python/kernel_tests/periodic_resample_op_test.py"], + srcs_version = "PY2AND3", + tags = [ + "notap", + ], + deps = [ + ":init_py", + "//tensorflow/contrib/util:util_py", + "//tensorflow/python:framework_test_lib", + ], +) + # py_library( # name = "periodic_resample_op_py", # srcs = ["python/ops/periodic_resample_op.py"], diff --git a/tensorflow/contrib/periodic_resample/kernels/periodic_resample_op.h b/tensorflow/contrib/periodic_resample/kernels/periodic_resample_op.h index bef21f7a5c..ba410f025d 100644 --- a/tensorflow/contrib/periodic_resample/kernels/periodic_resample_op.h +++ b/tensorflow/contrib/periodic_resample/kernels/periodic_resample_op.h @@ -100,6 +100,8 @@ template <class InputDataT, desired_shape.size(), ".")); bool found = false; + const auto& input_tensor_shape = input_tensor.shape(); + for (int i = 0; i < rank; ++i) { // if (desired_shape(i) < 1) { if (desired_shape[i] < 1) { @@ -111,6 +113,15 @@ template <class InputDataT, adjustable_dimension = i; found = true; } else { + OP_REQUIRES( + context, desired_shape[i] >= input_tensor_shape.dim_size(i), + tensorflow::errors::InvalidArgument( + "periodic_resample expects the size of non-adjustable " + "dimensions be at least as large as size of input tensor." + " Dimension ", i, " input tensor has size ", + input_tensor_shape.dim_size(i), ", desired shape has size ", + desired_shape[i], ".")); + // target_dimensions[i] = desired_shape(i); target_dimensions[i] = desired_shape[i]; new_sliced_size *= target_dimensions[i]; diff --git a/tensorflow/contrib/periodic_resample/ops/array_ops.cc b/tensorflow/contrib/periodic_resample/ops/array_ops.cc index c90fc06c7f..82bd796956 100644 --- a/tensorflow/contrib/periodic_resample/ops/array_ops.cc +++ b/tensorflow/contrib/periodic_resample/ops/array_ops.cc @@ -34,26 +34,40 @@ This function implements a slightly more generic version of the subpixel convolutions found in this [paper](https://arxiv.org/abs/1609.05158). The formula for computing the elements in the `output` tensor is as follows: + `T` = `values` tensor of rank `R` + `S` = desired `shape` of output tensor (vector of length `R`) + `P` = `output` tensor of rank `R` - \((T_1,\ldots,T_R)\) = shape(`T`) - \([S_1,\ldots,S_q,\ldots,S_R]\) = elements of vector `S` - A single element in `S` is left unspecified (denoted \(S_q=-1\)). - Let \(f_i\) denote the (possibly non-integer) factor that relates the original - dimension to the desired dimensions, \(S_i=f_i T_i\), for \(i\neq q\) where - \(f_i>0\). + \\((T_1,\\ldots,T_R)\\) = shape(`T`) + + \\([S_1,\\ldots,S_q,\\ldots,S_R]\\) = elements of vector `S` + + A single element in `S` is left unspecified (denoted \\(S_q=-1\\)). + + Let \\(f_i\\) denote the (possibly non-integer) factor that relates the original + dimension to the desired dimensions, \\(S_i=f_i T_i\\), for \\(i\\neq q\\) where + \\(f_i>0\\). + Define the following: - \(g_i=\lceil f_i\rceil\) - \(t=\prod_i T_i\) - \(s=\prod_{i\neq q} S_i\) - \(S_q\) can then be defined as by \(S_q=\lfloor t/s\rfloor\). + + \\(g_i=\\lceil f_i\\rceil\\) + + \\(t=\\prod_i T_i\\) + + \\(s=\\prod_{i\\neq q} S_i\\) + + \\(S_q\\) can then be defined by \\(S_q=\\lfloor t/s\\rfloor\\). The elements of the resulting tensor are defined as - \(P_{s_1,\ldots,s_R}=T_{h_1,\ldots,h_q,\ldots,h_R}\). - The \(h_i\) (\(i\neq q\)) are defined by \(h_i=\lfloor s_i/g_i\rfloor\). - \(h_q=S_q\sum_{j\neq q}^{q-1}G_j \mathrm{mod}(s_j,g_j) + s_q\), where - \(G_j=\prod_{i}^{j-1}g_i\) (\(G_0=1\)). + + \\(P_{s_1,\\ldots,s_R}=T_{h_1,\\ldots,h_q,\\ldots,h_R}\\). + + The \\(h_i\\) (\\(i\\neq q\\)) are defined by \\(h_i=\\lfloor s_i/g_i\\rfloor\\). + + \\(h_q=S_q\\sum_{j\\neq q}^{q-1}G_j \\mathrm{mod}(s_j,g_j) + s_q\\), where + \\(G_j=\\prod_{i}^{j-1}g_i\\) (\\(G_0=1\\)). One drawback of this method is that whenever the output dimensions are slightly less than integer multiples of the input dimensions, many of the tensor elements diff --git a/tensorflow/contrib/periodic_resample/python/kernel_tests/periodic_resample_op_test.py b/tensorflow/contrib/periodic_resample/python/kernel_tests/periodic_resample_op_test.py index 1d727870f6..30a2077570 100644 --- a/tensorflow/contrib/periodic_resample/python/kernel_tests/periodic_resample_op_test.py +++ b/tensorflow/contrib/periodic_resample/python/kernel_tests/periodic_resample_op_test.py @@ -19,8 +19,9 @@ from __future__ import division from __future__ import print_function import numpy -import tensorflow + from tensorflow.contrib.periodic_resample import periodic_resample +from tensorflow.python.framework import errors_impl from tensorflow.python.framework import test_util from tensorflow.python.ops import variables from tensorflow.python.platform import googletest @@ -96,6 +97,19 @@ class PeriodicResampleTest(test_util.TensorFlowTestCase): result = periodic_resample(input_tensor, desired_shape).eval() self.assertAllEqual(result, output_tensor) + def testPeriodicResampleErrors(self): + input_tensor = numpy.zeros(shape=[1, 2, 2, 4]) + with self.test_session(): + variables.global_variables_initializer().run() + with self.assertRaisesWithPredicateMatch( + errors_impl.InvalidArgumentError, + 'Dimension 3 input tensor has size 4, desired shape has size 1'): + periodic_resample(input_tensor, [None, 4, 4, 1]).eval() + with self.assertRaisesWithPredicateMatch( + errors_impl.InvalidArgumentError, + '4, to be the same as the length of the desired shape, 3'): + periodic_resample(input_tensor, [None, 4, 4]).eval() + if __name__ == "__main__": googletest.main() diff --git a/tensorflow/contrib/rnn/python/kernel_tests/core_rnn_cell_test.py b/tensorflow/contrib/rnn/python/kernel_tests/core_rnn_cell_test.py index b5d81b7caa..cafeb56ad8 100644 --- a/tensorflow/contrib/rnn/python/kernel_tests/core_rnn_cell_test.py +++ b/tensorflow/contrib/rnn/python/kernel_tests/core_rnn_cell_test.py @@ -663,6 +663,12 @@ class DropoutWrapperTest(test.TestCase): self.assertEqual(res[1].h.shape, (batch_size, 3)) return res + def testWrappedCellProperty(self): + cell = rnn_cell_impl.BasicRNNCell(10) + wrapper = rnn_cell_impl.DropoutWrapper(cell) + # Github issue 15810 + self.assertEqual(wrapper.wrapped_cell, cell) + def testDropoutWrapperKeepAllConstantInput(self): keep = array_ops.ones([]) res = self._testDropoutWrapper( diff --git a/tensorflow/contrib/rnn/python/kernel_tests/rnn_cell_test.py b/tensorflow/contrib/rnn/python/kernel_tests/rnn_cell_test.py index 73789206f3..70aaba1728 100644 --- a/tensorflow/contrib/rnn/python/kernel_tests/rnn_cell_test.py +++ b/tensorflow/contrib/rnn/python/kernel_tests/rnn_cell_test.py @@ -1549,5 +1549,100 @@ class BenchmarkLSTMCellXLA(test.Benchmark): benchmark_results["wall_time"]]])) +class WeightNormLSTMCellTest(test.TestCase): + """Compared cell output with pre-calculated values.""" + + def _cell_output(self, cell): + """Calculate cell output""" + + with self.test_session() as sess: + init = init_ops.constant_initializer(0.5) + with variable_scope.variable_scope("root", + initializer=init): + x = array_ops.zeros([1, 2]) + c0 = array_ops.zeros([1, 2]) + h0 = array_ops.zeros([1, 2]) + + state0 = rnn_cell.LSTMStateTuple(c0, h0) + + xout, sout = cell()(x, state0) + + sess.run([variables.global_variables_initializer()]) + res = sess.run([xout, sout], { + x.name: np.array([[1., 1.]]), + c0.name: 0.1 * np.asarray([[0, 1]]), + h0.name: 0.1 * np.asarray([[2, 3]]), + }) + + actual_state_c = res[1].c + actual_state_h = res[1].h + + return actual_state_c, actual_state_h + + def testBasicCell(self): + """Tests cell w/o peepholes and w/o normalisation""" + + def cell(): + return contrib_rnn_cell.WeightNormLSTMCell(2, + norm=False, + use_peepholes=False) + + actual_c, actual_h = self._cell_output(cell) + + expected_c = np.array([[0.65937078, 0.74983585]]) + expected_h = np.array([[0.44923624, 0.49362513]]) + + self.assertAllClose(expected_c, actual_c, 1e-5) + self.assertAllClose(expected_h, actual_h, 1e-5) + + def testNonbasicCell(self): + """Tests cell with peepholes and w/o normalisation""" + + def cell(): + return contrib_rnn_cell.WeightNormLSTMCell(2, + norm=False, + use_peepholes=True) + + actual_c, actual_h = self._cell_output(cell) + + expected_c = np.array([[0.65937084, 0.7574988]]) + expected_h = np.array([[0.4792085, 0.53470564]]) + + self.assertAllClose(expected_c, actual_c, 1e-5) + self.assertAllClose(expected_h, actual_h, 1e-5) + + + def testBasicCellWithNorm(self): + """Tests cell w/o peepholes and with normalisation""" + + def cell(): + return contrib_rnn_cell.WeightNormLSTMCell(2, + norm=True, + use_peepholes=False) + + actual_c, actual_h = self._cell_output(cell) + + expected_c = np.array([[0.50125383, 0.58805949]]) + expected_h = np.array([[0.32770363, 0.37397948]]) + + self.assertAllClose(expected_c, actual_c, 1e-5) + self.assertAllClose(expected_h, actual_h, 1e-5) + + def testNonBasicCellWithNorm(self): + """Tests cell with peepholes and with normalisation""" + + def cell(): + return contrib_rnn_cell.WeightNormLSTMCell(2, + norm=True, + use_peepholes=True) + + actual_c, actual_h = self._cell_output(cell) + + expected_c = np.array([[0.50125383, 0.59587258]]) + expected_h = np.array([[0.35041603, 0.40873795]]) + + self.assertAllClose(expected_c, actual_c, 1e-5) + self.assertAllClose(expected_h, actual_h, 1e-5) + if __name__ == "__main__": test.main() diff --git a/tensorflow/contrib/rnn/python/ops/rnn_cell.py b/tensorflow/contrib/rnn/python/ops/rnn_cell.py index e4667828cd..d7ae6621db 100644 --- a/tensorflow/contrib/rnn/python/ops/rnn_cell.py +++ b/tensorflow/contrib/rnn/python/ops/rnn_cell.py @@ -38,6 +38,7 @@ from tensorflow.python.ops import random_ops from tensorflow.python.ops import rnn_cell_impl from tensorflow.python.ops import variable_scope as vs from tensorflow.python.ops import partitioned_variables +from tensorflow.python.ops import nn_impl from tensorflow.python.platform import tf_logging as logging from tensorflow.python.util import nest @@ -328,7 +329,7 @@ class TimeFreqLSTMCell(rnn_cell_impl.RNNCell): def __init__(self, num_units, use_peepholes=False, cell_clip=None, initializer=None, num_unit_shards=1, forget_bias=1.0, - feature_size=None, frequency_skip=None, + feature_size=None, frequency_skip=1, reuse=None): """Initialize the parameters for an LSTM cell. @@ -2723,3 +2724,257 @@ class SRUCell(rnn_cell_impl._LayerRNNCell): h = r * self._activation(c) + (1.0 - r) * inputs return h, c + + +class WeightNormLSTMCell(rnn_cell_impl.RNNCell): + """Weight normalized LSTM Cell. Adapted from `rnn_cell_impl.LSTMCell`. + + The weight-norm implementation is based on: + https://arxiv.org/abs/1602.07868 + Tim Salimans, Diederik P. Kingma. + Weight Normalization: A Simple Reparameterization to Accelerate + Training of Deep Neural Networks + + The default LSTM implementation based on: + http://www.bioinf.jku.at/publications/older/2604.pdf + S. Hochreiter and J. Schmidhuber. + "Long Short-Term Memory". Neural Computation, 9(8):1735-1780, 1997. + + The class uses optional peephole connections, optional cell clipping + and an optional projection layer. + + The optional peephole implementation is based on: + https://research.google.com/pubs/archive/43905.pdf + Hasim Sak, Andrew Senior, and Francoise Beaufays. + "Long short-term memory recurrent neural network architectures for + large scale acoustic modeling." INTERSPEECH, 2014. + """ + + def __init__(self, num_units, norm=True, use_peepholes=False, + cell_clip=None, initializer=None, num_proj=None, + proj_clip=None, forget_bias=1, activation=None, + reuse=None): + """Initialize the parameters of a weight-normalized LSTM cell. + + Args: + num_units: int, The number of units in the LSTM cell + norm: If `True`, apply normalization to the weight matrices. If False, + the result is identical to that obtained from `rnn_cell_impl.LSTMCell` + use_peepholes: bool, set `True` to enable diagonal/peephole connections. + cell_clip: (optional) A float value, if provided the cell state is clipped + by this value prior to the cell output activation. + initializer: (optional) The initializer to use for the weight matrices. + num_proj: (optional) int, The output dimensionality for the projection + matrices. If None, no projection is performed. + proj_clip: (optional) A float value. If `num_proj > 0` and `proj_clip` is + provided, then the projected values are clipped elementwise to within + `[-proj_clip, proj_clip]`. + forget_bias: Biases of the forget gate are initialized by default to 1 + in order to reduce the scale of forgetting at the beginning of + the training. + activation: Activation function of the inner states. Default: `tanh`. + reuse: (optional) Python boolean describing whether to reuse variables + in an existing scope. If not `True`, and the existing scope already has + the given variables, an error is raised. + """ + super(WeightNormLSTMCell, self).__init__(_reuse=reuse) + + self._scope = 'wn_lstm_cell' + self._num_units = num_units + self._norm = norm + self._initializer = initializer + self._use_peepholes = use_peepholes + self._cell_clip = cell_clip + self._num_proj = num_proj + self._proj_clip = proj_clip + self._activation = activation or math_ops.tanh + self._forget_bias = forget_bias + + self._weights_variable_name = "kernel" + self._bias_variable_name = "bias" + + if num_proj: + self._state_size = rnn_cell_impl.LSTMStateTuple(num_units, num_proj) + self._output_size = num_proj + else: + self._state_size = rnn_cell_impl.LSTMStateTuple(num_units, num_units) + self._output_size = num_units + + @property + def state_size(self): + return self._state_size + + @property + def output_size(self): + return self._output_size + + def _normalize(self, weight, name): + """Apply weight normalization. + + Args: + weight: a 2D tensor with known number of columns. + name: string, variable name for the normalizer. + Returns: + A tensor with the same shape as `weight`. + """ + + output_size = weight.get_shape().as_list()[1] + g = vs.get_variable(name, [output_size], dtype=weight.dtype) + return nn_impl.l2_normalize(weight, dim=0) * g + + def _linear(self, args, + output_size, + norm, + bias, + bias_initializer=None, + kernel_initializer=None): + """Linear map: sum_i(args[i] * W[i]), where W[i] is a variable. + + Args: + args: a 2D Tensor or a list of 2D, batch x n, Tensors. + output_size: int, second dimension of W[i]. + bias: boolean, whether to add a bias term or not. + bias_initializer: starting value to initialize the bias + (default is all zeros). + kernel_initializer: starting value to initialize the weight. + + Returns: + A 2D Tensor with shape [batch x output_size] equal to + sum_i(args[i] * W[i]), where W[i]s are newly created matrices. + + Raises: + ValueError: if some of the arguments has unspecified or wrong shape. + """ + if args is None or (nest.is_sequence(args) and not args): + raise ValueError("`args` must be specified") + if not nest.is_sequence(args): + args = [args] + + # Calculate the total size of arguments on dimension 1. + total_arg_size = 0 + shapes = [a.get_shape() for a in args] + for shape in shapes: + if shape.ndims != 2: + raise ValueError("linear is expecting 2D arguments: %s" % shapes) + if shape[1].value is None: + raise ValueError("linear expects shape[1] to be provided for shape %s, " + "but saw %s" % (shape, shape[1])) + else: + total_arg_size += shape[1].value + + dtype = [a.dtype for a in args][0] + + # Now the computation. + scope = vs.get_variable_scope() + with vs.variable_scope(scope) as outer_scope: + weights = vs.get_variable( + self._weights_variable_name, [total_arg_size, output_size], + dtype=dtype, + initializer=kernel_initializer) + if norm: + wn = [] + st = 0 + with ops.control_dependencies(None): + for i in range(len(args)): + en = st + shapes[i][1].value + wn.append(self._normalize(weights[st:en, :], + name='norm_{}'.format(i))) + st = en + + weights = array_ops.concat(wn, axis=0) + + if len(args) == 1: + res = math_ops.matmul(args[0], weights) + else: + res = math_ops.matmul(array_ops.concat(args, 1), weights) + if not bias: + return res + + with vs.variable_scope(outer_scope) as inner_scope: + inner_scope.set_partitioner(None) + if bias_initializer is None: + bias_initializer = init_ops.constant_initializer(0.0, dtype=dtype) + + biases = vs.get_variable( + self._bias_variable_name, [output_size], + dtype=dtype, + initializer=bias_initializer) + + return nn_ops.bias_add(res, biases) + + def call(self, inputs, state): + """Run one step of LSTM. + + Args: + inputs: input Tensor, 2D, batch x num_units. + state: A tuple of state Tensors, both `2-D`, with column sizes + `c_state` and `m_state`. + + Returns: + A tuple containing: + + - A `2-D, [batch x output_dim]`, Tensor representing the output of the + LSTM after reading `inputs` when previous state was `state`. + Here output_dim is: + num_proj if num_proj was set, + num_units otherwise. + - Tensor(s) representing the new state of LSTM after reading `inputs` when + the previous state was `state`. Same type and shape(s) as `state`. + + Raises: + ValueError: If input size cannot be inferred from inputs via + static shape inference. + """ + dtype = inputs.dtype + num_units = self._num_units + sigmoid = math_ops.sigmoid + c, h = state + + input_size = inputs.get_shape().with_rank(2)[1] + if input_size.value is None: + raise ValueError("Could not infer input size from inputs.get_shape()[-1]") + + with vs.variable_scope(self._scope, initializer=self._initializer): + + concat = self._linear([inputs, h], 4 * num_units, + norm=self._norm, bias=True) + + # i = input_gate, j = new_input, f = forget_gate, o = output_gate + i, j, f, o = array_ops.split(value=concat, num_or_size_splits=4, axis=1) + + if self._use_peepholes: + w_f_diag = vs.get_variable("w_f_diag", shape=[num_units], dtype=dtype) + w_i_diag = vs.get_variable("w_i_diag", shape=[num_units], dtype=dtype) + w_o_diag = vs.get_variable("w_o_diag", shape=[num_units], dtype=dtype) + + new_c = (c * sigmoid(f + self._forget_bias + w_f_diag * c) + + sigmoid(i + w_i_diag * c) * self._activation(j)) + else: + new_c = (c * sigmoid(f + self._forget_bias) + + sigmoid(i) * self._activation(j)) + + if self._cell_clip is not None: + # pylint: disable=invalid-unary-operand-type + new_c = clip_ops.clip_by_value(new_c, -self._cell_clip, self._cell_clip) + # pylint: enable=invalid-unary-operand-type + if self._use_peepholes: + new_h = sigmoid(o + w_o_diag * new_c) * self._activation(new_c) + else: + new_h = sigmoid(o) * self._activation(new_c) + + if self._num_proj is not None: + with vs.variable_scope("projection"): + new_h = self._linear(new_h, + self._num_proj, + norm=self._norm, + bias=False) + + if self._proj_clip is not None: + # pylint: disable=invalid-unary-operand-type + new_h = clip_ops.clip_by_value(new_h, + -self._proj_clip, + self._proj_clip) + # pylint: enable=invalid-unary-operand-type + + new_state = rnn_cell_impl.LSTMStateTuple(new_c, new_h) + return new_h, new_state diff --git a/tensorflow/contrib/seq2seq/python/kernel_tests/beam_search_decoder_test.py b/tensorflow/contrib/seq2seq/python/kernel_tests/beam_search_decoder_test.py index d2beac5f31..f498b2bb57 100644 --- a/tensorflow/contrib/seq2seq/python/kernel_tests/beam_search_decoder_test.py +++ b/tensorflow/contrib/seq2seq/python/kernel_tests/beam_search_decoder_test.py @@ -225,6 +225,94 @@ class TestBeamStep(test.TestCase): self.assertAllEqual(next_state_.log_probs, expected_log_probs) +class TestLargeBeamStep(test.TestCase): + """ + Tests a single step of beam search in such + case that beam size is larger than vocabulary size. + """ + + def setUp(self): + super(TestLargeBeamStep, self).setUp() + self.batch_size = 2 + self.beam_width = 8 + self.vocab_size = 5 + self.end_token = 0 + self.length_penalty_weight = 0.6 + + + def test_step(self): + def get_probs(): + """this simulates the initialize method in BeamSearchDecoder""" + log_prob_mask = array_ops.one_hot(array_ops.zeros([self.batch_size], + dtype=dtypes.int32), + depth=self.beam_width, on_value=True, + off_value=False, dtype=dtypes.bool) + + log_prob_zeros = array_ops.zeros([self.batch_size, self.beam_width], + dtype=dtypes.float32) + log_prob_neg_inf = array_ops.ones([self.batch_size, self.beam_width], + dtype=dtypes.float32) * -np.Inf + + log_probs = array_ops.where(log_prob_mask, log_prob_zeros, + log_prob_neg_inf) + return log_probs + + log_probs = get_probs() + dummy_cell_state = array_ops.zeros([self.batch_size, self.beam_width]) + + _finished = array_ops.one_hot( + array_ops.zeros([self.batch_size], dtype=dtypes.int32), + depth=self.beam_width, on_value=False, + off_value=True, dtype=dtypes.bool) + _lengths = np.zeros([self.batch_size, self.beam_width], dtype=np.int64) + _lengths[:, 0]=2 + _lengths = constant_op.constant(_lengths, dtype=dtypes.int64) + + beam_state = beam_search_decoder.BeamSearchDecoderState( + cell_state=dummy_cell_state, + log_probs=log_probs, + lengths=_lengths, + finished=_finished) + + logits_ = np.full([self.batch_size, self.beam_width, self.vocab_size], + 0.0001) + logits_[0, 0, 2] = 1.9 + logits_[0, 0, 3] = 2.1 + logits_[0, 1, 3] = 3.1 + logits_[0, 1, 4] = 0.9 + logits_[1, 0, 1] = 0.5 + logits_[1, 1, 2] = 2.7 + logits_[1, 2, 2] = 10.0 + logits_[1, 2, 3] = 0.2 + logits = constant_op.constant(logits_, dtype=dtypes.float32) + log_probs = nn_ops.log_softmax(logits) + + outputs, next_beam_state = beam_search_decoder._beam_search_step( + time=2, + logits=logits, + next_cell_state=dummy_cell_state, + beam_state=beam_state, + batch_size=ops.convert_to_tensor(self.batch_size), + beam_width=self.beam_width, + end_token=self.end_token, + length_penalty_weight=self.length_penalty_weight) + + with self.test_session() as sess: + outputs_, next_state_, state_, log_probs_ = sess.run( + [outputs, next_beam_state, beam_state, log_probs]) + + self.assertEqual(outputs_.predicted_ids[0, 0], 3) + self.assertEqual(outputs_.predicted_ids[0, 1], 2) + self.assertEqual(outputs_.predicted_ids[1, 0], 1) + neg_inf = -np.Inf + self.assertAllEqual(next_state_.log_probs[:, -3:], + [[neg_inf, neg_inf, neg_inf], + [neg_inf, neg_inf, neg_inf]]) + self.assertEqual((next_state_.log_probs[:, :-3] > neg_inf).all(), True) + self.assertEqual((next_state_.lengths[:, :-3] > 0).all(), True) + self.assertAllEqual(next_state_.lengths[:, -3:], [[0, 0, 0], + [0, 0, 0]]) + class BeamSearchDecoderTest(test.TestCase): def _testDynamicDecodeRNN(self, time_major, has_attention): diff --git a/tensorflow/contrib/seq2seq/python/ops/beam_search_decoder.py b/tensorflow/contrib/seq2seq/python/ops/beam_search_decoder.py index ebe25ce077..a5f7169c31 100644 --- a/tensorflow/contrib/seq2seq/python/ops/beam_search_decoder.py +++ b/tensorflow/contrib/seq2seq/python/ops/beam_search_decoder.py @@ -19,7 +19,6 @@ from __future__ import division from __future__ import print_function import collections - import numpy as np from tensorflow.contrib.seq2seq.python.ops import beam_search_ops @@ -229,8 +228,11 @@ class BeamSearchDecoder(decoder.Decoder): self._start_tokens = array_ops.tile( array_ops.expand_dims(self._start_tokens, 1), [1, self._beam_width]) self._start_inputs = self._embedding_fn(self._start_tokens) - self._finished = array_ops.zeros( - [self._batch_size, self._beam_width], dtype=dtypes.bool) + + self._finished = array_ops.one_hot( + array_ops.zeros([self._batch_size], dtype=dtypes.int32), + depth=self._beam_width, on_value=False, + off_value=True, dtype=dtypes.bool) @property def batch_size(self): @@ -298,11 +300,15 @@ class BeamSearchDecoder(decoder.Decoder): """ finished, start_inputs = self._finished, self._start_inputs + log_probs = array_ops.one_hot( # shape(batch_sz, beam_sz) + array_ops.zeros([self._batch_size], dtype=dtypes.int32), + depth=self._beam_width, on_value=0.0, off_value=-np.Inf, + dtype=nest.flatten(self._initial_cell_state)[0].dtype) + + initial_state = BeamSearchDecoderState( cell_state=self._initial_cell_state, - log_probs=array_ops.zeros( - [self._batch_size, self._beam_width], - dtype=nest.flatten(self._initial_cell_state)[0].dtype), + log_probs=log_probs, finished=finished, lengths=array_ops.zeros( [self._batch_size, self._beam_width], dtype=dtypes.int64)) @@ -563,18 +569,11 @@ def _beam_search_step(time, logits, next_cell_state, beam_state, batch_size, time = ops.convert_to_tensor(time, name="time") # During the first time step we only consider the initial beam scores_shape = array_ops.shape(scores) - scores_flat = control_flow_ops.cond( - time > 0, - lambda: array_ops.reshape(scores, [batch_size, -1]), - lambda: scores[:, 0]) - num_available_beam = control_flow_ops.cond( - time > 0, lambda: math_ops.reduce_prod(scores_shape[1:]), - lambda: math_ops.reduce_prod(scores_shape[2:])) + scores_flat = array_ops.reshape(scores, [batch_size, -1]) # Pick the next beams according to the specified successors function - next_beam_size = math_ops.minimum( - ops.convert_to_tensor(beam_width, dtype=dtypes.int32, name="beam_width"), - num_available_beam) + next_beam_size = ops.convert_to_tensor(beam_width, dtype=dtypes.int32, + name="beam_width") next_beam_scores, word_indices = nn_ops.top_k(scores_flat, k=next_beam_size) next_beam_scores.set_shape([static_batch_size, beam_width]) diff --git a/tensorflow/contrib/verbs/BUILD b/tensorflow/contrib/verbs/BUILD index 38a84ffb10..80a5d07ea4 100644 --- a/tensorflow/contrib/verbs/BUILD +++ b/tensorflow/contrib/verbs/BUILD @@ -99,7 +99,7 @@ cc_library( alwayslink = 1, ) -tf_cuda_library( +cc_library( name = "rdma_rendezvous_mgr", srcs = ["rdma_rendezvous_mgr.cc"], hdrs = ["rdma_rendezvous_mgr.h"], @@ -114,7 +114,7 @@ tf_cuda_library( ], ) -cc_library( +tf_cuda_library( name = "rdma_mgr", srcs = ["rdma_mgr.cc"], hdrs = ["rdma_mgr.h"], @@ -141,6 +141,8 @@ tf_cuda_library( "//conditions:default": [], }), deps = [ + ":grpc_verbs_client", + ":verbs_service_proto_cc", ":verbs_util", "//tensorflow/core:core_cpu_internal", "//tensorflow/core:framework", diff --git a/tensorflow/contrib/verbs/README.md b/tensorflow/contrib/verbs/README.md index 7c1c8ea459..1b99f4ce4f 100644 --- a/tensorflow/contrib/verbs/README.md +++ b/tensorflow/contrib/verbs/README.md @@ -24,66 +24,144 @@ The design is based on TensorFlow r1.0. An RDMA path is added between servers fo During the server setup, an RDMA manager is created to manage low-level RDMA components such as RDMA channel and RDMA adapter, an RDMA rendezvous manager is created to oversee send/recv operations between servers. Following the distributed TensorFlow design philosophy, the send operation is passive, i.e. merely placing a tensor in the local out-going table. It is the receive operation that actually initiates the tensor transfer. -TensorFlow dynamically allocates memory for tensors that are to be sent or received. This causes difficulty for RDMA operations where pinned memory is required. Two remedies are possible, either the memory is pinned, transfer, then unpinned for each and every tensor to be transferred, or a buffer is pre-allocated and pinned for each tensor. The former incurs significant operation overhead since pinning and unpinning memory for each dynamically generated tensor is slow. The latter incurs large memory overhead and extra copying from the tensor to its pinned buffer, but may still be faster than the former. The second approach is adopted in this design. Each RDMA channel, representing a RDMA connection to a peer, contains a table of pinned buffers for all the seen tensors that requires transfer. It is assumed that the tensor size rarely changes across different steps. So only one buffer is created for the same tensor across all the steps. In the rare case when the tensor size does increases, the old buffer is discarded and new buffer of larger size is created and pinned. +TensorFlow dynamically allocates memory for tensors that are to be sent or received. This causes difficulty for RDMA operations where pinned memory is required. Few remedies are possible: +1. The memory is pinned, transfered, then unpinned for each and every tensor to be transferred. This incurs significant operation overhead since pinning and unpinning memory for each dynamically generated tensor is slow. +2. Buffer is pre-allocated and pinned for each tensor. This incurs large memory overhead and extra copying from the tensor to its pinned buffer, but may still be faster than the former. +3. Following HKUST research on the use of GPU direct, and their [GDR implementation](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/contrib/gdr/README.md), there is a smart way to benefit from the TensorFlow allocation theme which is mostly pool based, i.e allocators pre-allocate a large memory block, and allocate the tensors from there. By attaching a custom Visitor to relevant alloactors, we can do a single registration of the entire memory block, which zeros the registration overhead. Once the block is registered, each new tensor allocated will be at a registred address, which will allow us to do direct RDMA writes to it. -When a tensor is prepared for transfer, it is first converted to TensorProto, then the proto is serialized to byte array and copied to the pinned buffer. The content of the buffer is transferred to the remote node via RDMA write. On the remote side, the process is reversed. This is illustrated in the diagram below. The conversion of TensorProto is introduced to simplify transfer of string-tensors. Also since the TensorProto lives in host memory, even if the origin tensor lives in the device, the pinned buffers are all allocated in the host memory. - +For best performance, we will adopt HKUST 0 copies approach in our solution. This means: + +1. Tensor writes will be done directly from the source tensor to the **result** tensor, with no memory copies in between. This should be done for all DMAable tensors which are located either on CPU or on a RDMA compatible GPU device (GPU direct). +2. Non DMAable tensors (CanMemCopy == false) will be serialized to a TensorProto on the sender side, RDMA written to a registered buffer on the receiver side, and then deserialized by the receiver. +3. Tensors which are located on a non-RDMA-compatible GPU, will be RDMA written to a registered CPU **proxy** buffer on the receiver side, and then copied to GPU by the receiver. -The following improvements can be made in the future. First, conversion to TensorProto and serialization can be avoided for numeric (float/int) tensors since their internal buffer can be access directly as byte array. Second, the pinned buffer may be allocated on device if the tensor is located in the device. This avoids extra device-to-host copy at the expense of extra device memory consumption. ## Design details -### RDMA components +### Terminology -* **RDMA adapter:** The base for RDMA communications. It may contain multiple channels and buffers. It is responsible for handling various incoming RDMA messages. -* **RDMA channel:** Responsible for RDMA connection to a particular node. It manages multiple buffers. A channel has a callback table which stores all the callbacks for the requested tensors. -* **RDMA buffer:** Responsible for sending or receiving data. It has a fixed size memory to store the data. It has a queue to store the pending jobs. There are three types of buffers, message buffer, ACK buffer and tensor buffer. A channel has two message buffers, two ack buffers and many tensor buffers. -* **RDMA manager:** Manages the adapter and channels, including channel creation, channel setup via GRPC service, channel lookup, etc. -* **RDMA rendezvous manager:** manages multiple rdma rendezvous. -* **RDMA rendezvous:** a derived class of BaseRemoteRendezvous. This class is the back end for "send" and "recv" ops. When the sendrecv_op wants to send or receive a tensor, it calls the rendezvous' "send" and "recv" functions respectively. Rendezvous are identified by "step_id", a random number, so that tensors for different iterations don't get mixed up. +* **Sender** - The node which sends the tensor. +* **Receiver** - The node which receives the tensor. +* **Result tensor** - The destination tensor, allocated on its appropriate device. +* **Proxy tensor** - A CPU allocated tensor, which will be used in the case where the result tensor cannot be RDMA written to directly (GPU direct is disabled or not available). The RDMA write will therefore be done to the proxy tensor, and afterwards we will do a manual local copy from it to the result tensor. -### The SEND operation +### Messages -In TensorFlow, when rendezvous sends a tensor, it merely puts a tensor in a local table in the corresponding rendezvous. If the tensor has been requested, a callback exists in the table. "send" will activate the callback, which tries to send the tensor across the node. +* RDMA_MESSAGE_TENSOR_REQUEST +* RDMA_MESSAGE_META_DATA_RESPONSE +* RDMA_MESSAGE_TENSOR_RE_REQUEST +### Transport protocol -### The RECV operation +The tensor transfer process is initiated when the receiver requests a tensor. In code it is done by calling **Rendezvous::Recv()** or **Rendezvous::RecvAsync()**. The TensorFlow base implementation handles the case where the requested tensor is located on the same node. The more interesting case where the requested tensor is located on a remote node (receiver != sender) is to be handled in a derivation of the pure virtual **BaseRemoteRendezvous::RecvFromRemoteAsync()**. TensorFlow provides a default GRPC based implementation which comes in the vanilla version but suffers in scalability when running large models. Our RDMA based implementation presumes to be more scalable. HKUST's contrib GDR implementation is more scalable than GRPC, and less scalable than ours, only because we did our evolution based on it. -When a tensor is requested, rendezvous' recv function is called. The function first places a callback in the channel's callback table, which will be activated once the tensor is sent from the source. In the next step, a message is sent to notify the source of the requested tensor. Once the source receives the message, it will check locally for the tensor, if not found, a callback is placed in the table, otherwise, the tensor id will be placed at corresponding RDMA buffer's job queue for future transmission. When a tensor is scheduled to be transmitted, the RDMA buffer needs to have the memory allocated and initialized (registered with the remote buffer info). If the memory is not ready, the transmission is deferred, a message is sent to the destination to establish the memory first. The other case a transmission can be deferred is when the buffer is still being used by an on-going transmission. +Our entry point is the implementation of **RdmaRemoteRendezvous::RecvFromRemoteAsync()**, located in rdma_rendezvous_mgr.cc. The implementation creates a new **RdmaTensorRequest** object, keyed by request index (uint32_t), stores it in a list of pending requests, and calls its **Start()** method. The **Start()** method basically does 2 things: -### Three types of RDMA buffers +1. Allocate the result tensor (and the proxy tensor if required). +2. Send a **RDMA_MESSAGE_TENSOR_REQUEST** to the sender, containing the address of the destination tensor (result/proxy) for RDMA write. -* **Message buffer:** responsible for sending message only. -* **Ack buffer:** once a message is sent, the recipient needs to send an ack via the ack buffer to free up the message buffer. An ack buffer is exclusively for its coupled message buffer. -* **Tensor buffer:** responsible for sending tensors. The recipient needs to send back a message to free up the sending buffer. +In order to allocate the result and proxy tensors, we need to know the tensor's meta-data, i.e. shape and data-type for DMAable tensors, and proto-size for serialized tensors. Unfortunately, this information is only available on the sender side which complicates manners. In order to avoid sending extra messages for querying the meta-data at each step, we store a local meta-data cache per tensor, which will only be update upon changes. Based on the assumption that the meta-data of a tensor rarely changes between steps, we expect that on most times the cache will only be updated once. The sender is responsible to detect changes in the meta-data, and update the receiver. In order for the sender to know that the meta-data had changed, each **RDMA_MESSAGE_TENSOR_REQUEST** will contain the meta-data that the receiver had grabbed from the local cache. The sender will then compare the meta-data from the message to the tensor's new meta-data. -### RDMA packet format +When the sender receives an **RDMA_MESSAGE_TENSOR_REQUEST**, it will create a new **RdmaTensorResponse** object for the given request message, store it in a list of pending responses, and will invoke its **Start()** method. The **Start()** method does the following: -|type|name_size|name|step_id|buffer_size|remote_addr|rkey|is_dead|data_type|tensor_shape|tensor_bytes|tensor_buffer| +1. Grab the source tensor from the local table (In code, **RecvLocalAsync()**). +2. If the source tensor is not DMAable, serialize it to a TensorProto. +3. If the source tensor is located on a device which cannot be DMA written from, copy it to CPU. +4. If it is the first time this tensor is requested, or if the tensor's meta-data changed: + 1. Clone the tensor's data to be sent later. + 2. Send a **RDMA_MESSAGE_META_DATA_RESPONSE** containing the new meta-data. +5. Otherwise: + 1. RDMA write the tensor (or TensorProto) to the destination address and rkey specified in the request message. The immediate value for the write will be the request index. -### Six types of RDMA messages -* RDMA_MESSAGE_ACK -* RDMA_MESSAGE_BUFFER_IDLE -* RDMA_MESSAGE_BUFFER_REQUEST -* RDMA_MESSAGE_BUFFER_RESPONSE -* RDMA_MESSAGE_TENSOR_REQUEST -* RDMA_MESSAGE_TENSOR_WRITE - -### Actions upon receiving RDMA messages -* RDMA_MESSAGE_ACK - * sender: mark local ack buffer idle. - * receiver: mark remote message buffer idle, send next item. -* RDMA_MESSAGE_BUFFER_IDLE - * sender: mark local message buffer idle, send next item. - * receiver: send ack, set remote tensor buffer idle, send next item. -* RDMA_MESSAGE_BUFFER_REQUEST - * sender: mark local message buffer idle, send next item. - * receiver: send ack, find or create tensor buffer, send BUFFER_RESPONSE. -* RDMA_MESSAGE_BUFFER_RESPONSE - * sender: mark local message buffer idle, send next item. - * receiver: send ack, set remote buffer info, set local and remote buffer idle, send next item. -* RDMA_MESSAGE_TENSOR_REQUEST - * sender: mark local message buffer idle, send next item. - * receiver: send ack, find or create tensor buffer, enqueue tensor id, send next item. -* RDMA_MESSAGE_TENSOR_WRITE - * sender: mark local message buffer idle, send next item. - * receiver: run callback. + +When the receiver receives the **RDMA_MESSAGE_META_DATA_RESPONSE**, it will locate the relevant **RdmaTensorRequest** using the request index specified in the message, and invoke its **RecvTensorMetaData()** which does the following: + +1. Update the local meta-data cache. +2. Reallocate the result/proxy tensors. +3. Re-send the tensor request. For tracability, the new message has a different name: **RDMA_MESSAGE_TENSOR_RE_REQUEST**. + +When the sender receives a **RDMA_MESSAGE_TENSOR_RE_REQUEST**, it will locate the relevant **RdmaTensorResponse** using the request index specified in the message, and invoke its **Resume()** method, which will RDMA write the contents of the tensor that was cloned earlier, to the new remote address specified in the re-request. + +When the receiver receives the RDMA write, it will locate the relevant **RdmaTensorRequest** using the request index which is the immediate value. It will then invoke its **RecvTensorContent()** which does the following: + +1. Proxy copy/deserialize if required. +2. Invoke the done callback. +3. Deallocate the result/proxy tensors and remove the request from the pending list. + + + +### Additional design notes + +1. When the sender receives a tensor request, the source tensor may or may not be ready yet. The situation is handled through a process of tag matching: + * If the request arrives before the tensor is ready, then a callback is put in a local table, and will be invoked once the tensor arrives. + * If the tensor is ready before the request arives, than the tensor is put in a local table. When the request arrives, it will invoke the callback immediatly. + In code it is done by calling **RecvLocalAsync()**, which receives the tensor's key, step-id, and the callback. +2. When the callback is invoked, the relevant tensor is removed from the tag matching table. In the case where we need to send the tensor's meta-data, the **RdmaTensorResponse** will store a copy of the tensor until the re-request arrives. +3. The sending of protocol messages (**RDMA_MESSAGE_TENSOR_REQUEST**, **RDMA_MESSAGE_META_DATA_RESPONSE** and **RDMA_MESSAGE_TENSOR_RE_REQUEST**) is done by the class **RdmaMessageBuffer**. All messages are sent using RDMA writes from/to fixed messages buffers. This implies that we cannot send on a specific channel more than one message at a time. In order to synchronize the messages, the **RdmaMessageBuffer** holds the a local and remote buffer statuses which can be either busy or idle. When a write is issued, both statuses will be changed to busy. When the write-complete event is received, the local status is changed to idle. When the write is received on the remote side, the remote side will parse the message, and return an ACK back to the sending side on which the sending side will update the remote status to idle. When both the local and remote statuses are idle, the next message can be sent. +5. ACK writes are empty writes (hence they require no buffer) with immediate value 0xFFFFFFFE. Message writes have the immediate value 0xFFFFFFFF. All other writes are tensor-content writes whose immediate value is the request-index. + +### RDMA components + +* **enum RdmaImmDataType** - Immediate types to distinguish between different RDMA writes on the remote side. Ack writes and control-message writes have a fixed immediate value. The rest of the writes are tensor writes and the immediate value is the relevant request index. +* **enum RdmaWriteIDType** - Types to distinguish between different RDMA write-complete events: Ack, control message and tensor writes. +* **class RdmaWriteID** - Context for RDMA write complete events. Holds the RdmaWriteIDType and additional data. +* **class RdmaTensorMetaData** - Meta-data for a tensor (type, shape, is_dead, proto_size). +* **class RdmaMemoryMgr** - Manages the meta-data cache, and the registered memory regions. +* **class RdmaTensorRequest** - Holds and manages information for a single tensor request throughout the entire receive cycle. API: + * **Start()** - Start the request sequence. + * Allocate the result tensor (and proxy tensor if required). + * Send RDMA_MESSAGE_TENSOR_REQUEST to the remote side. + * **RecvTensorMetaData()** - Receive meta-data from the remote side. + * Update the local meta-data cache. + * Reallocate the result tensor (and proxy tensor if required). + * Re-send the request to the remote side. + * **RecvTensorContent()** - Receive tensor content from the remote side (RDMA write was completed). + * Decode proto if required and/or move to GPU if the content was not written to it directly (GPU direct is not avaliable). + * Invoke the done callback. +* **class RdmaTensorResponse** - Holds and manages information for a single tensor response throughout the entire send cycle. API: + * **Start()** - Start the response sequence. + * Find the tensor in the local tag-match table. + * Compare the tensor's meta-data to the meta-data in the message (taken from the requester's local cache). + * If meta-data changed: + * Clone the tensor to be sent later. + * Send a meta-data update message and wait for re-request. + * Else: + * Send the tensor's content (using direct RDMA write). + * **Resume()** - Resume the response sequence after a re-request. Send the tensor's content that was cloned earlier. + * **Destroy()** - Destroy the response's resources and remove it form the pending list. +* **class RdmaAdapter** - The base for RDMA communications. It may contain multiple channels and buffers. It is responsible for handling various incoming RDMA messages. +* **class RdmaChannel** - Responsible for RDMA connection to a particular node. It manages messagee buffers. A channel has a request table which stores all the pending tensor requests. +* **class RdmaMessageBuffer** - Responsible for sending or receiving messages. It has a fixed size memory to store the data. It has a queue to store the pending jobs. A channel has two message buffers one for tx and one for rx. +* **class RdmaMgr** - Manages the adapter and channels, including channel creation, channel setup via GRPC service, channel lookup, etc. +* **class RdmaRendezvousMgr** - Manages multiple rdma rendezvous. +* **class RdmaRemoteRendezvous** - A derived class of BaseRemoteRendezvous. This class is the back end for "send" and "recv" ops. When the sendrecv_op wants to send or receive a tensor, it calls the rendezvous' "send" and "recv" functions respectively. Rendezvous are identified by "step_id", a random number, so that tensors for different iterations don't get mixed up. + +### Message structure: + +| type | name_size | name | step_id | request_index | remote_addr/checksum | rkey | is_dead | data_type | tensor_shape | tensor_bytes | error_status | +|------|---------- |------|---------|---------------|----------------------|------|---------|-----------|--------------|--------------|-----------------------| +| 1B | 2B | 512 | 8B | 8B | 8B | 4B | 1B | XB | XB | 8B | Size - 4B, proto - XB | + +* **RDMA_MESSAGE_TENSOR_REQUEST** - (receiver ==> sender) The original tensor request. + * type - The message type. + * name (name_size) - Name of the requested tensor. + * step_id - Step ID. + * request_index - Request index. + * remote_addr/rkey - Address/rkey of the result/proxy tensor. Irrelevant for first-time request. + * is_dead/data_type/tensor_shape/tensor_bytes - The current meta-data as stored in the receiver local cache. The sender will use that information to know if the receiver's cache requires updating. +* **RDMA_MESSAGE_META_DATA_RESPONSE** - (sender ==> receiver) The meta-data update message in case meta-data had changed (or if it is the first time the tensor is requested). + * type - The message type. + * request_index - Request index. + * is_dead/data_type/tensor_shape/tensor_bytes - The up-to-date meta-data. + * checksum - In data validation mode, this will hold the checksum of the source tensor. +* **RDMA_MESSAGE_TENSOR_RE_REQUEST** - (receiver ==> sender) Tensor re-requset after meta-data update and reallocation of result/proxy tensors. + * type - The message type. + * name (name_size) - Name of the requested tensor. + * step_id - Step ID. + * request_index - Request index. + * remote_addr/rkey - Address/rkey of the reallocated result/proxy tensor. +* **RDMA_MESSAGE_ERROR_STATUS** - (sender ==> receiver) Notify the receiver that an error had occured on the sender side, so it can propagate it to the upper levels. + * type - The message type. + * name (name_size) - Name of the requested tensor. + * step_id - Step ID. + * request_index - Request index. + * error_status - The error status (code, message, details). diff --git a/tensorflow/contrib/verbs/grpc_verbs_service.cc b/tensorflow/contrib/verbs/grpc_verbs_service.cc index f2af6b79fb..742f946c95 100644 --- a/tensorflow/contrib/verbs/grpc_verbs_service.cc +++ b/tensorflow/contrib/verbs/grpc_verbs_service.cc @@ -122,17 +122,15 @@ Status GrpcVerbsService::GetRemoteAddressSync( rc->SetRemoteAddress(ra, false); rc->Connect(); int i = 0; - int idx[] = {1, 0, 3, 2}; - std::vector<RdmaBuffer*> mb(rc->message_buffers()); - CHECK_EQ(request->mr_size(), 4); + int idx[] = {1, 0}; + std::vector<RdmaMessageBuffer*> mb(rc->message_buffers()); + CHECK_EQ(request->mr_size(), RdmaChannel::kNumMessageBuffers); for (const auto& mr : request->mr()) { // the connections are crossed, i.e. // local tx_message_buffer <---> remote rx_message_buffer_ // local rx_message_buffer <---> remote tx_message_buffer_ - // local tx_ack_buffer <---> remote rx_ack_buffer_ - // local rx_ack_buffer <---> remote tx_ack_buffer_ - // hence idx[] = {1, 0, 3, 2}. - RdmaBuffer* rb = mb[idx[i]]; + // hence idx[] = {1, 0}. + RdmaMessageBuffer* rb = mb[idx[i]]; RemoteMR rmr; rmr.remote_addr = mr.remote_addr(); rmr.rkey = mr.rkey(); diff --git a/tensorflow/contrib/verbs/patch_notes_verbs_with_0_copies.md b/tensorflow/contrib/verbs/patch_notes_verbs_with_0_copies.md new file mode 100644 index 0000000000..956b8f2147 --- /dev/null +++ b/tensorflow/contrib/verbs/patch_notes_verbs_with_0_copies.md @@ -0,0 +1,87 @@ +## Verbs implementation to use direct tensor writes (0 copies) + +### Motivation: + +Following HKUST research on the use of GPU direct, and their [GDR implementation](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/contrib/gdr/README.md), we wish to adopt the 0 copies approach and apply it to the current verbs implementation, while keeping the current implementation advantages, such as configurability and the use of RDMA for control messages. + +### Performance: + +Compared with the current GRPC, verbs and GDR implementation, the result implementation gave the best performance for every model, with any number of nodes. For VGG16 on 8 nodes with 4 P100 GPUs each, the prototype beat the second place by over 15%. + +### Implementation requirements: + +1. Tensor writes need to be done directly from the source Tensor to the destination Tensor, with no memory copies in between. This should be done for all DMAble tensors which are located either on CPU or on a RDMA compatible GPU device (GPU direct). +2. Non DMAble tensors (CanMemCopy == false) will be serialized to proto on the sender side, RDMA written to a registered buffer on the receiver side, and then deserialized by the receiver. +3. Tensors which are located on a non-RDMA-compatible GPU, will be RDMA written to a registered CPU proxy buffer on the receiver side, and then copied to GPU by the receiver. + +### Implementation constrains: + +For best stability and proof of correctness, we will divide the implementation to two stages: +1. At first stage we will keep changes to the current implementation to the minimum possible. The expense will be that we may have unused or unnecessary code leftovers, which may also affect performance. +2. At second stage, we will re-iterate over the code and remove irrelevant code parts. +The design of the solution aims that we will achieve both stages with relative ease. + +### Design guidelines: + +1. Since we do not want to do any unnecessary memory copying, we will no longer allocate a fixed CPU buffer as the destination for the RDMA write. Instead we will do the writing directly to the result tensor, or if the result tensor is on a device which does not support RDMA, we will do the writing to a proxy CPU tensor and then copy its content to the result tensor. +2. The address of the destination Tensor needs to be sent to the sender side for writing, meaning that the result/proxy tensor should be pre-allocated on the receiver side, prior to sending the tensor request. In order to do that, we need to know its meta-data, i.e. shape and data-type for DMAble tensors, and proto-size for serialized tensors. Unfortunately, this information is only available on the sender side which complicates manners. In order to avoid sending extra messages for querying the meta-data on each step, we store a local meta-data cache per tensor. Based on the assumption that the meta-data of a tensor rarely changes between steps, we expect that on most times the cache will only be updated once. When the sender receives a request for a tensor, if it is the first time this tensor is requested, or in the rare case that the meta-data did change, the sender will first send a meta-data response, on which the receiver will update the local cache, and reallocate the result/proxy tensors if required. When the receiver sends the tensor request, it will contain also the meta-data currently stored in its local cache, so the sender can compare it to see if there was a change. +3. When the sender writes the tensor content to the result tensor, no additional data is being written with it. That means we need to reside on ibverbs immediate (uint32_t) to indicate which request we are responding to (in order to trigger the receive callback). The easiest and most elegant way is to key the recv callback with a unique request_index (uint32_t), instead of the current key_with_step_id (string). +4. Since the sender no longer writes the tensor from/to fixed buffers, we no longer need to schedule the writes using the local/remote status. In addition we no longer rely on the RmdaTensorBuffer members as the source/destination addresses and rkey/lkey. Instead, each RdmaTensorBuffer will hold multiple "Response" objects (one per step-id), from which we derive destination address and rkey. The source address and lkey are always the ones of the source Tensor. +5. With the addition of tensor pre-allocation, we noticed there is a large code similarity between sending the first tensor request and re-sending the request in case of meta-data changes. After implementing a common method for tensor pre-allocation, it turned out that implementation becomes much simpler by encapsulating the process of request sending/re-sending, meta-data response callback and content response callback, all in a single "Request" class. The request class holds all the relevant request information, which reduces excessive parameter passing and lambda capturing. This decision is purely for elegance and code simplicity, and we decided to implement it in first stage because it makes the implementation much easier. + +### New types/classes: + +* **enum RdmaImmDataType** - Immediate types to distinguish between different RDMA writes on the remote side. Ack writes and control-message writes have a fixed immediate value. The rest of the writes are tensor writes and the immediate value is the relevant request index. +* **enum RdmaWriteIDType** - Types to distinguish between different RDMA write-complete events: Ack, control message, tensor DMA write and tensor proto write. +* **class RdmaWriteID** - Context for RDMA write complete events. Holds the RdmaWriteIDType and additional data. +* **class RemoteAddressContext** - Remote address information (address + mr). Will be passed as write context for tensor proto writes. +* **class RdmaTensorMetaData** - Meta-data for a tensor (type, shape, is_dead, proto_size). +* **class RdmaMemoryMgr** - Manages the meta-data cache, and the registered memory regions. +* **class RdmaTensorRequest** - Holds and manages information for a single tensor request throughout the entire receive cycle. API: + * Start() - Start the request. + * RecvTensorMetaData() - Receive meta-data from the remote side. + * RecvTensorContent() - Receive tensor content from the remote side and invoke the done() callback. +* **class RdmaTensorResponse** - Holds information for a single tensor response, such as destination address and rkey. + +### Protocol changes: + +The protocol messages themselves will remain mostly unchanged at the first stage, but will be used differently, as described below. The current messages structures already have most of the required fields for the new implementation. The only change is the "buffer_size" field which is no longer used since we are no longer sending additional information with the tensor, and thus it is now always equal to the "tensor_bytes" field. Instead, we use that field to pass the "request_index". + +### Message structure: + +| type | name_size | name | step_id | request_index | remote_addr | rkey | is_dead | data_type | tensor_shape | tensor_bytes | +|------|---------- |------|---------|---------------|-------------|------|---------|-----------|--------------|--------------| +| 1B | 2B | 512 | 8B | 8B | 8B | 4B | 1B | XB | XB | 8B | + +* **RDMA_MESSAGE_TENSOR_REQUEST** - (receiver ==> sender) The original tensor request. + * type - The message type. + * name (name_size) - Name of the requested tensor. + * step_id - Step ID. + * request_index - Request index. + * remote_addr/rkey - Address/rkey of the result/proxy tensor. Irrelevant for first-time request. + * is_dead/data_type/tensor_shape/tensor_bytes - The current meta-data as stored in the receiver local cache. The sender will use that information to know if the receiver's cache requires updating. +* **RDMA_MESSAGE_BUFFER_REQUEST** - (sender ==> receiver) The meta-data update message in case meta-data had changed (or if it is the first time the tensor is requested). + * type - The message type. + * request_index - Request index. + * is_dead/data_type/tensor_shape/tensor_bytes - The up-to-date meta-data. +* **RDMA_MESSAGE_BUFFER_RESPONSE** - (receiver ==> sender) Tensor re-requset after meta-data update and reallocation of result/proxy tensors. + * type - The message type. + * name (name_size) - Name of the requested tensor. + * step_id - Step ID. + * request_index - Request index. + * remote_addr/rkey - Address/rkey of the reallocated result/proxy tensor. + * is_dead/data_type/tensor_shape/tensor_bytes - The new meta-data. Will be removed in the next phase. +* **RDMA_MESSAGE_TENSOR_WRITE** - (sender ==> receiver) No longer sent. There is only a direct write of the tensor content to the result/proxy tensor. Request index passed as the immediate value of the write. +* **RDMA_MESSAGE_TENSOR_IDLE** - (receiver ==> sender) No longer sent. + + + +### Second stage optimizations: +1. Remove unused code leftovers. +2. Remove the ACK buffer completely, since we can rely completely on its immediate value. + +### Future optimizations: +1. Map the tensor names to indexes, to significantly reduce the request message size. +2. Understand the purpose of empty tensors and if we can skip remote fetching for them. +3. Consider concatenating multiple requests and/or using multiple message buffers. +4. Consider a no-request architecture. diff --git a/tensorflow/contrib/verbs/rdma.cc b/tensorflow/contrib/verbs/rdma.cc index ae9a384565..ec5271abe0 100644 --- a/tensorflow/contrib/verbs/rdma.cc +++ b/tensorflow/contrib/verbs/rdma.cc @@ -16,18 +16,19 @@ limitations under the License. #ifdef TENSORFLOW_USE_VERBS #include "tensorflow/contrib/verbs/rdma.h" -#include <fcntl.h> +#include "tensorflow/contrib/verbs/verbs_service.pb.h" #include <cstdlib> #include <fcntl.h> -#include "tensorflow/contrib/verbs/verbs_util.h" #include "tensorflow/core/common_runtime/device_mgr.h" #include "tensorflow/core/common_runtime/dma_helper.h" +#include "tensorflow/core/common_runtime/process_util.h" #if GOOGLE_CUDA #include "tensorflow/core/common_runtime/gpu/gpu_util.h" #include "tensorflow/core/common_runtime/gpu/process_state.h" #endif #include "tensorflow/core/distributed_runtime/rendezvous_mgr_interface.h" #include "tensorflow/core/distributed_runtime/session_mgr.h" +#include "tensorflow/core/distributed_runtime/rpc/grpc_util.h" #include "tensorflow/core/framework/rendezvous.h" #include "tensorflow/core/framework/tensor.h" #include "tensorflow/core/lib/core/status.h" @@ -41,32 +42,19 @@ namespace tensorflow { #define RoCE_V2 "RoCE v2" namespace { -// hash name to 32-bit integer -uint32_t NameHash(const string& name) { - return Hash32(name.data(), name.size(), 0x1234ABCD); -} // convenience function for printing message string MessageTypeToString(RdmaMessageType rmt) { switch (rmt) { - case RDMA_MESSAGE_ACK: - return "RDMA_MESSAGE_ACK"; - break; - case RDMA_MESSAGE_BUFFER_IDLE: - return "RDMA_MESSAGE_BUFFER_IDLE"; - break; - case RDMA_MESSAGE_BUFFER_REQUEST: - return "RDMA_MESSAGE_BUFFER_REQUEST"; + case RDMA_MESSAGE_META_DATA_UPDATE: + return "RDMA_MESSAGE_META_DATA_UPDATE"; break; - case RDMA_MESSAGE_BUFFER_RESPONSE: - return "RDMA_MESSAGE_BUFFER_RESPONSE"; + case RDMA_MESSAGE_TENSOR_RE_REQUEST: + return "RDMA_MESSAGE_TENSOR_RE_REQUEST"; break; case RDMA_MESSAGE_TENSOR_REQUEST: return "RDMA_MESSAGE_TENSOR_REQUEST"; break; - case RDMA_MESSAGE_TENSOR_WRITE: - return "RDMA_MESSAGE_TENSOR_WRITE"; - break; default: return "UNKNOWN MESSAGE"; } @@ -347,7 +335,7 @@ uint32_t set_param(uint32_t default_val, const char* env_param) { enum ibv_mtu set_mtu(uint8_t port_num, ibv_context* context) { ibv_port_attr port_attr; - enum ibv_mtu mtu; + enum ibv_mtu mtu = IBV_MTU_512; string mtu_s; int rc, mtu_i; @@ -468,97 +456,70 @@ void RdmaAdapter::Process_CQ() { rc->Recv(); // imm_data is the index of RX buffer in the buffer table. uint32_t imm_data = wc_[i].imm_data; - RdmaBuffer* rb = rc->FindBuffer(imm_data); + RdmaMessageBuffer* rb; RdmaMessage rm; - RdmaMessage::ParseMessage(rm, rb->buffer_); - VLOG(2) << "recv RDMA message: " << MessageTypeToString(rm.type_); - if (rm.type_ == RDMA_MESSAGE_ACK) { + if (imm_data == RDMA_IMM_DATA_ACK) { // receive an ack to a message rb = rc->tx_message_buffer_; rb->SetBufferStatus(remote, idle); rb->SendNextItem(); - } else if (rm.type_ == RDMA_MESSAGE_TENSOR_REQUEST) { - // received a request-for-tensor message - // send ack to release remote tx message buffer - RdmaBuffer* ab = rc->tx_ack_buffer_; - ab->SendNextItem(); - // find or create buffer - RdmaBuffer* tb = rc->FindOrCreateBuffer(rm.name_); - string key_with_step_id = - VerbsUtil::AppendStepidToKey(rm.name_, rm.step_id_); - tb->EnqueueItem(key_with_step_id); - // send the next tensor - worker_env_->compute_pool->Schedule([tb]() { tb->SendNextItem(); }); - } else if (rm.type_ == RDMA_MESSAGE_BUFFER_IDLE) { - // receive tensor-buffer-ready message - // send ack to release remote tx message buffer - RdmaBuffer* ab = rc->tx_ack_buffer_; - ab->SendNextItem(); - // find buffer - RdmaTensorBuffer* tb = - reinterpret_cast<RdmaTensorBuffer*>(rc->FindBuffer(rm.name_)); - tb->SetBufferStatus(remote, idle); - worker_env_->compute_pool->Schedule([tb]() { tb->ReSendNextItem(); }); - } else if (rm.type_ == RDMA_MESSAGE_BUFFER_REQUEST) { - // remote host requests to create a tensor buffer; - // send ack to release remote tx message buffer - RdmaBuffer* ab = rc->tx_ack_buffer_; - ab->SendNextItem(); - // find or create the buffer - RdmaBuffer* tb = rc->FindOrCreateBuffer(rm.name_, TENSOR); - RemoteMR rmr; - rmr.remote_addr = rm.remote_addr_; - rmr.rkey = rm.rkey_; - tb->SetRemoteMR(rmr, true); - tb->CreateCPUBuffer(rm.buffer_size_); - // create RDMA_MESSAGE_BUFFER_RESPONSE message - RdmaMessage br; - br.type_ = RDMA_MESSAGE_BUFFER_RESPONSE; - br.name_size_ = rm.name_.size(); - br.name_ = rm.name_; - br.buffer_size_ = rm.buffer_size_; - br.remote_addr_ = reinterpret_cast<uint64_t>(tb->buffer_); - br.rkey_ = tb->self_->rkey; - string message = RdmaMessage::CreateMessage(br); - RdmaBuffer* mb = rc->tx_message_buffer_; - mb->EnqueueItem(message); - mb->SendNextItem(); - } else if (rm.type_ == RDMA_MESSAGE_BUFFER_RESPONSE) { - // remote creates a buffer and responds - // send ack to release remote tx message buffer - RdmaBuffer* ab = rc->tx_ack_buffer_; - ab->SendNextItem(); - // find buffer - RdmaTensorBuffer* tb = - reinterpret_cast<RdmaTensorBuffer*>(rc->FindBuffer(rm.name_)); - CHECK(rm.buffer_size_ == tb->size_) - << "rm.buffer_size = " << rm.buffer_size_ - << "tb->size_ = " << tb->size_ << "rm.name_ = " << rm.name_; - RemoteMR rmr; - rmr.remote_addr = rm.remote_addr_; - rmr.rkey = rm.rkey_; - tb->SetRemoteMR(rmr, true); - tb->SetBufferStatus(local, idle); - tb->SetBufferStatus(remote, idle); - worker_env_->compute_pool->Schedule([tb]() { tb->ReSendNextItem(); }); - } else if (rm.type_ == RDMA_MESSAGE_TENSOR_WRITE) { - // tensor RDMA write completed - worker_env_->compute_pool->Schedule([rm, rc]() { - string key_with_step_id = - VerbsUtil::AppendStepidToKey(rm.name_, rm.step_id_); - rc->RunRecvCallback(key_with_step_id); - }); + continue; } - } else if (wc_[i].opcode == IBV_WC_RDMA_WRITE) { - RdmaBuffer* rb = reinterpret_cast<RdmaBuffer*>(wc_[i].wr_id); - rb->SetBufferStatus(local, idle); - RdmaMessage rm; + + if (imm_data <= RDMA_IMM_MAX_REQUEST_ID) { + // receive a tensor RDMA write + uint32_t request_index = imm_data; + RdmaTensorRequest* request = rc->GetTensorRequest(request_index); + request->RecvTensorContent(); + continue; + } + + // receive a control message + rb = rc->rx_message_buffer_; RdmaMessage::ParseMessage(rm, rb->buffer_); - VLOG(2) << "sent RDMA message: " << MessageTypeToString(rm.type_); - if (rm.type_ != RDMA_MESSAGE_ACK) { - worker_env_->compute_pool->Schedule([rb]() { rb->SendNextItem(); }); + RdmaMessageBuffer::SendAck(rc); + RDMA_LOG(1) << "Step 0x" << std::hex << rm.step_id_ << std::dec + << ": Received " << MessageTypeToString(rm.type_) << " " + << "#" << rm.request_index_ << ": " << rm.name_; + + if (rm.type_ == RDMA_MESSAGE_TENSOR_REQUEST) { + RdmaTensorResponse* response = rc->AddTensorResponse(rm); + response->Start(); + } else if (rm.type_ == RDMA_MESSAGE_META_DATA_UPDATE) { + RdmaTensorRequest* request = rc->GetTensorRequest(rm.request_index_); + request->RecvTensorMetaData(rm.data_type_, rm.tensor_shape_, + rm.is_dead_, rm.tensor_bytes_); +#ifdef RDMA_DATA_VALIDATION + request->RecvTensorChecksum(rm.checksum_); +#endif + } else if (rm.type_ == RDMA_MESSAGE_TENSOR_RE_REQUEST) { + RdmaTensorResponse* response = rc->UpdateTensorResponse(rm); + response->Resume(); + } else if (rm.type_ == RDMA_MESSAGE_ERROR_STATUS) { + RdmaTensorRequest* request = rc->GetTensorRequest(rm.request_index_); + request->RecvErrorStatus(rm.status_); + } + } else if (wc_[i].opcode == IBV_WC_RDMA_WRITE) { + RdmaWriteID* wr_id = reinterpret_cast<RdmaWriteID*>(wc_[i].wr_id); + RDMA_LOG(2) << "Write complete of type " << wr_id->write_type; + switch (wr_id->write_type) { + case RDMA_WRITE_ID_ACK: + break; + case RDMA_WRITE_ID_MESSAGE: { + RdmaMessageBuffer* rb = + reinterpret_cast<RdmaMessageBuffer*>(wr_id->write_context); + rb->SetBufferStatus(local, idle); + rb->SendNextItem(); + break; + } + case RDMA_WRITE_ID_TENSOR_WRITE: { + RdmaTensorResponse* response = + reinterpret_cast<RdmaTensorResponse*>(wr_id->write_context); + response->Destroy(); + } } + delete wr_id; } } } @@ -588,8 +549,10 @@ int RdmaChannel::PingPostSend() { RdmaChannel::RdmaChannel(const RdmaAdapter* adapter, const string local_name, const string remote_name) - : adapter_(adapter), local_name_(local_name), remote_name_(remote_name) { - + : adapter_(adapter), + local_name_(local_name), + remote_name_(remote_name), + request_serial_(0) { struct ibv_sge list; mr_ = ibv_reg_mr(adapter_->pd_, ping_buff_, kPingBuffSize, @@ -651,29 +614,15 @@ RdmaChannel::RdmaChannel(const RdmaAdapter* adapter, const string local_name, // create message and ack buffers, then initialize the tables. { - const string buffer_names[] = {"tx_message_buffer", "rx_message_buffer", - "tx_ack_buffer", "rx_ack_buffer"}; + const string buffer_names[] = {"tx_message_buffer", "rx_message_buffer"}; tx_message_buffer_ = new RdmaMessageBuffer(this, buffer_names[0]); rx_message_buffer_ = new RdmaMessageBuffer(this, buffer_names[1]); - tx_ack_buffer_ = new RdmaAckBuffer(this, buffer_names[2]); - rx_ack_buffer_ = new RdmaAckBuffer(this, buffer_names[3]); message_buffers_.reserve(kNumMessageBuffers); message_buffers_.push_back(tx_message_buffer_); message_buffers_.push_back(rx_message_buffer_); - message_buffers_.push_back(tx_ack_buffer_); - message_buffers_.push_back(rx_ack_buffer_); // create buffer on host tx_message_buffer_->CreateCPUBuffer(RdmaMessage::kRdmaMessageBufferSize); rx_message_buffer_->CreateCPUBuffer(RdmaMessage::kRdmaMessageBufferSize); - tx_ack_buffer_->CreateCPUBuffer(RdmaMessage::kRdmaAckBufferSize); - rx_ack_buffer_->CreateCPUBuffer(RdmaMessage::kRdmaAckBufferSize); - // bt_mu_.lock() is not used in constructor. - for (int i = 0; i < kNumMessageBuffers; i++) { - uint32_t index = NameHash(buffer_names[i]); - buffer_table_.insert({index, message_buffers_[i]}); - buffer_index_name_table_.insert({index, buffer_names[i]}); - buffer_name_index_table_.insert({buffer_names[i], index}); - } } CHECK(PingPostRecv() == 0) << "Couldn't post receive from " << remote_name_ << " with error " << std::strerror(errno); @@ -684,8 +633,6 @@ RdmaChannel::~RdmaChannel() { CHECK(!ibv_destroy_qp(qp_)) << "Failed to destroy QP"; delete tx_message_buffer_; delete rx_message_buffer_; - delete tx_ack_buffer_; - delete rx_ack_buffer_; } void RdmaChannel::SetRemoteAddress(const RdmaAddress& ra, bool override) { @@ -716,114 +663,31 @@ void RdmaChannel::Recv() { CHECK(!ibv_post_recv(qp_, &wr, &bad_wr)) << "Failed to post recv"; } -// Lookup 32-bit buffer index from buffer name -// Args: -// buffer_name: name of the buffer -// Returns: -// 32-bit index -uint32_t RdmaChannel::LookupBufferIndex(const string& buffer_name) { - mutex_lock lock{bt_mu_}; - BufferNameIndexTable::iterator iter = - buffer_name_index_table_.find(buffer_name); - CHECK(iter != buffer_name_index_table_.end()); - return iter->second; -} - -// Find a buffer by its 32-bit index -// Args: -// index: 32-bit hash code of the tensor buffer name -// Returns: -// name of the tensor buffer -RdmaBuffer* RdmaChannel::FindBuffer(const uint32_t index) { - mutex_lock lock{bt_mu_}; - BufferTable::iterator iter = buffer_table_.find(index); - CHECK(iter != buffer_table_.end()); - return iter->second; -} - -// Find a buffer by its name -// Args: -// name: name of the buffer -// Returns: -// the named rdma buffer -RdmaBuffer* RdmaChannel::FindBuffer(const string& name) { - uint32_t index = LookupBufferIndex(name); - return FindBuffer(index); -} - -// Find a buffer if it exists, otherwise create one. -// The memory inside the created buffer is not allocated. -// Args: -// name: the name of the buffer -// buffer_type: TENSOR, MESSAGE or ACK. -// Returns: -// the named buffer -RdmaBuffer* RdmaChannel::FindOrCreateBuffer(const string& name, - BufferType buffer_type) { - mutex_lock lock{bt_mu_}; - RdmaBuffer* rb; - // find index - BufferNameIndexTable::iterator iter = buffer_name_index_table_.find(name); - if (iter != buffer_name_index_table_.end()) { - uint32_t index = iter->second; - // find buffer - BufferTable::iterator iter = buffer_table_.find(index); - CHECK(iter != buffer_table_.end()); - rb = iter->second; - } else { - uint32_t index = NameHash(name); - if (buffer_type == TENSOR) { - rb = new RdmaTensorBuffer(this, name); - } else if (buffer_type == MESSAGE) { - rb = new RdmaMessageBuffer(this, name); - } else if (buffer_type == ACK) { - rb = new RdmaAckBuffer(this, name); - } - buffer_name_index_table_.insert({name, index}); - buffer_index_name_table_.insert({index, name}); - buffer_table_.insert({index, rb}); +RdmaTensorRequest* RdmaChannel::InsertTensorRequest( + const string& key, int64 step_id, Device* dst_dev, + const Rendezvous::Args recv_args, + const RdmaTensorRequest::RecvDoneCallback& done) { + mutex_lock lock{ct_mu_}; + uint32_t request_index = request_serial_++; + if (request_serial_ > RDMA_IMM_MAX_REQUEST_ID) { + request_serial_ = 0; } - CHECK(rb); - return rb; + RdmaTensorRequest request(request_index, key, step_id, this, dst_dev, + recv_args, done); + auto it = request_table_.emplace(request_index, request); + return &it.first->second; } -// Insert callback to the callback_table. -// The callback is activated when the corresponding tensor is received. -// Arg: -// key: the name of the tensor -// recv_done: the callback associated with the tensor. -// Returns: -// None -void RdmaChannel::InsertRecvCallback(const string& key, - std::function<void()> recv_done) { +void RdmaChannel::RemoveTensorRequest(uint32_t request_index) { mutex_lock lock{ct_mu_}; - callback_table_.insert({key, recv_done}); + request_table_.erase(request_index); } -// Remove callback from the callback_table. -// Arg: -// key: the name of the tensor -// Returns: -// None -void RdmaChannel::RemoveRecvCallback(const string& key) { +RdmaTensorRequest* RdmaChannel::GetTensorRequest(uint32_t request_index) { mutex_lock lock{ct_mu_}; - callback_table_.erase(key); -} - -// Run named callback in the callback_table. -// Arg: -// key: the name of the tensor -// Returns: -// None -void RdmaChannel::RunRecvCallback(const string& key) { - std::function<void()> recv_done; - { - mutex_lock lock{ct_mu_}; - CallbackTable::iterator iter = callback_table_.find(key); - CHECK(iter != callback_table_.end()); - recv_done = iter->second; - } - recv_done(); + RequestTable::iterator iter = request_table_.find(request_index); + CHECK(iter != request_table_.end()); + return &iter->second; } void RdmaChannel::Connect() { @@ -888,25 +752,22 @@ void RdmaChannel::Connect(const RdmaAddress& remoteAddr) { connected_ = true; } else { - LOG(INFO) << "channel already connected"; + RDMA_LOG(2) << "channel already connected"; } } -RdmaBuffer::RdmaBuffer(RdmaChannel* channel, string name) +RdmaMessageBuffer::RdmaMessageBuffer(RdmaChannel* channel, string name) : channel_(channel), name_(name) {} -RdmaBuffer::~RdmaBuffer() { +RdmaMessageBuffer::~RdmaMessageBuffer() { CHECK(!ibv_dereg_mr(self_)) << "ibv_dereg_mr failed"; FreeBuffer(); } -void RdmaBuffer::FreeBuffer() { +void RdmaMessageBuffer::FreeBuffer() { if ((buffer_ != nullptr) && buffer_on_host_) { free(buffer_); } - // TODO - // release buffer if it is on device. - // We don't support RDMABuffer on device at this moment. } // Allocate CPU memory for the Rdma buffer @@ -915,7 +776,7 @@ void RdmaBuffer::FreeBuffer() { // lock: whether or not mutex_lock the process to protect concurrency. // Returns: // None -void RdmaBuffer::CreateCPUBuffer(size_t size, bool lock) { +void RdmaMessageBuffer::CreateCPUBuffer(size_t size, bool lock) { CHECK(size > 0); if (lock) { mu_.lock(); @@ -943,7 +804,7 @@ void RdmaBuffer::CreateCPUBuffer(size_t size, bool lock) { // override: whether override existing information // Returns: // None -void RdmaBuffer::SetRemoteMR(RemoteMR rmr, bool override) { +void RdmaMessageBuffer::SetRemoteMR(RemoteMR rmr, bool override) { mutex_lock lock{mu_}; if ((override) || (remote_status_ == none)) { remote_.remote_addr = rmr.remote_addr; @@ -956,63 +817,51 @@ void RdmaBuffer::SetRemoteMR(RemoteMR rmr, bool override) { } // Put a task in the buffer's job queue -void RdmaBuffer::EnqueueItem(string item) { +void RdmaMessageBuffer::EnqueueItem(string item) { mutex_lock lock{mu_}; queue_.push(item); } // Rdma-Write the content of the buffer -void RdmaBuffer::Write(uint32_t imm_data, size_t buffer_size) { +void RdmaMessageBuffer::Write(uint32_t imm_data, size_t buffer_size) { + Write(channel_, imm_data, buffer_size, (uint64_t)buffer_, self_->lkey, + remote_.remote_addr, remote_.rkey, RDMA_WRITE_ID_MESSAGE, this); +} + +// Generalized Write method +void RdmaMessageBuffer::Write(const RdmaChannel* channel, uint32_t imm_data, + size_t buffer_size, uint64_t src_addr, + uint32_t lkey, uint64_t remote_addr, + uint32_t rkey, RdmaWriteIDType write_type, + void* write_context) { struct ibv_sge list; - list.addr = (uint64_t)buffer_; + list.addr = src_addr; list.length = buffer_size; - list.lkey = self_->lkey; + list.lkey = lkey; struct ibv_send_wr wr; memset(&wr, 0, sizeof(wr)); - wr.wr_id = (uint64_t) this; + wr.wr_id = (uint64_t) new RdmaWriteID(write_type, write_context); wr.sg_list = &list; wr.num_sge = 1; wr.opcode = IBV_WR_RDMA_WRITE_WITH_IMM; wr.send_flags = IBV_SEND_SIGNALED; wr.imm_data = imm_data; - wr.wr.rdma.remote_addr = (uint64_t)remote_.remote_addr; - wr.wr.rdma.rkey = remote_.rkey; + wr.wr.rdma.remote_addr = remote_addr; + wr.wr.rdma.rkey = rkey; struct ibv_send_wr* bad_wr; - CHECK(!ibv_post_send(channel_->qp_, &wr, &bad_wr)) << "Failed to post send"; -} - -RdmaAckBuffer::RdmaAckBuffer(RdmaChannel* channel, string name) - : RdmaBuffer(channel, name) {} - -RdmaMessageBuffer::RdmaMessageBuffer(RdmaChannel* channel, string name) - : RdmaBuffer(channel, name) {} - -RdmaTensorBuffer::RdmaTensorBuffer(RdmaChannel* channel, string name) - : RdmaBuffer(channel, name) {} - -RdmaTensorBuffer::~RdmaTensorBuffer() { - for (Itable it = retable.begin(); it != retable.end(); ++it) { - delete (it->second); - } + CHECK(!ibv_post_send(channel->qp_, &wr, &bad_wr)) << "Failed to post send"; } // Send the next ack from the buffer's job queue. -void RdmaAckBuffer::SendNextItem() { - uint32_t imm_data = LookupBufferIndex("rx_ack_buffer"); - RdmaMessage rm; - rm.name_ = "rx_ack_buffer"; - rm.type_ = RDMA_MESSAGE_ACK; - rm.name_size_ = rm.name_.size(); - string message = RdmaMessage::CreateMessage(rm); - memcpy(buffer_, message.data(), message.size()); - Write(imm_data, message.size()); +void RdmaMessageBuffer::SendAck(const RdmaChannel* channel) { + Write(channel, RDMA_IMM_DATA_ACK, 0, 0, 0, 0, 0, RDMA_WRITE_ID_ACK, nullptr); } // Send the next message from the buffer's job queue. void RdmaMessageBuffer::SendNextItem() { - uint32_t imm_data = LookupBufferIndex("rx_message_buffer"); + uint32_t imm_data = RDMA_IMM_DATA_MESSAGE; mu_.lock(); if (!queue_.empty() && (local_status_ == idle) && (remote_status_ == idle)) { local_status_ = busy; @@ -1029,244 +878,392 @@ void RdmaMessageBuffer::SendNextItem() { } } -Rendezvous::DoneCallback RdmaTensorBuffer::getRecvTensorCallback( - const string& key_with_step_id, const string& key, int64 step_id, - const Rendezvous::ParsedKey& parsed) { - Rendezvous::DoneCallback cb = [this, key_with_step_id, key, step_id, parsed]( - const Status& status, const Rendezvous::Args& send_args, - const Rendezvous::Args& recv_args, const Tensor& in, bool is_dead) { - CHECK(status.ok()) << "RecvLocalAsync was not ok, key" << key_with_step_id - << " error message: " << status.error_message(); - size_t buffer_size = RdmaMessage::kMessageTotalBytes; - size_t tensor_bytes = 0; - // Figures out which device the tensor is hosted on. - Device* src_dev = nullptr; - Status s = channel_->adapter_->worker_env_->device_mgr->LookupDevice( - parsed.src_device, &src_dev); - CHECK(s.ok()) << "src device not found"; - // Does the device have the right incarnation number we expect? - CHECK(src_dev->attributes().incarnation() == parsed.src_incarnation) - << "RecvTensor expects a different device incarnation: " - << parsed.src_incarnation << " vs. " - << src_dev->attributes().incarnation() - << ". Your worker job was probably restarted. Check your " - << "worker job for the reason why it was restarted."; - Device* dst_dev = nullptr; - // destination is on CPU. - s = channel_->adapter_->worker_env_->device_mgr->LookupDevice("CPU:0", - &dst_dev); - CHECK(s.ok()) << "dst device not found"; - AllocatorAttributes dst_alloc_attr; - dst_alloc_attr.set_on_host(true); - - bool can_memcpy = DataTypeCanUseMemcpy(in.dtype()); - // string tensor needs to be serialized - Tensor copy; - TensorProto proto; - if (src_dev->tensorflow_gpu_device_info() && - (!send_args.alloc_attrs.on_host())) { #if GOOGLE_CUDA - CHECK(send_args.device_context) << "send dev name: " << src_dev->name() - << " gpu_info: " - << src_dev->tensorflow_gpu_device_info(); - - if (can_memcpy) { - AllocatorAttributes host_alloc_attrs; - host_alloc_attrs.set_gpu_compatible(true); - host_alloc_attrs.set_on_host(true); - Allocator* alloc = ProcessState::singleton()->GetCUDAHostAllocator(0); - copy = Tensor(alloc, in.dtype(), in.shape()); - tensor_bytes = in.TotalBytes(); - buffer_size += tensor_bytes; - GPUUtil::CopyGPUTensorToCPU( - src_dev, send_args.device_context, &in, ©, - [this, copy, tensor_bytes, buffer_size, key, in, step_id, - key_with_step_id, is_dead, send_args, recv_args](const Status& s) { - CHECK(s.ok()) << "copy tensor from gpu sync"; - StringPiece copy_buf; - copy_buf = copy.tensor_data(); - PostCopyOperations(true, buffer_size, tensor_bytes, key, in, - step_id, is_dead, key_with_step_id, ©, - NULL, ©_buf, send_args, recv_args); - }); - } else { - // "val" is on a GPU. No longer uses GPUUtil to fill the proto, use - // aync instead - GPUUtil::SetProtoFromGPU( - in, src_dev, send_args.device_context, &proto, is_dead, - [this, proto, buffer_size, key, in, step_id, key_with_step_id, - is_dead, send_args, recv_args](const Status& s) mutable { - CHECK(s.ok()) << "copy proto from gpu sync"; - auto tensor_bytes = proto.ByteSize(); - buffer_size += tensor_bytes; - PostCopyOperations(false, buffer_size, tensor_bytes, key, in, - step_id, is_dead, key_with_step_id, NULL, - &proto, NULL, send_args, recv_args); - }); - } -#endif // GOOGLE_CUDA - } else { - // tensor is in CPU memory. - StringPiece copy_buf; - if (can_memcpy) { - copy_buf = in.tensor_data(); - tensor_bytes = in.TotalBytes(); - } else { - in.AsProtoTensorContent(&proto); - tensor_bytes = proto.ByteSize(); - } - buffer_size += tensor_bytes; - PostCopyOperations(can_memcpy, buffer_size, tensor_bytes, key, in, - step_id, is_dead, key_with_step_id, ©, &proto, - ©_buf, send_args, recv_args); +static void CountCopies(const std::string& key, void* src_addr, void* dst_addr, + size_t tensor_bytes, bool is_gpu_to_cpu) { +#ifdef RDMA_COUNT_COPIES + static uint64_t numGPUToCPUCopies = 0; + static uint64_t numGPUToCPUCopiedBytes = 0; + static uint64_t numCPUToGPUCopies = 0; + static uint64_t numCPUToGPUCopiedBytes = 0; + static uint64_t numTotalCopies = 0; + + if (is_gpu_to_cpu) { + ++numGPUToCPUCopies; + numGPUToCPUCopiedBytes += tensor_bytes; + } else { + ++numCPUToGPUCopies; + numCPUToGPUCopiedBytes += tensor_bytes; + } + if ((++numTotalCopies % 0x400) == 0) { + RDMA_LOG(0) << "Tensor copies:" + << " GPU to CPU: " << numGPUToCPUCopies + << " (" << numGPUToCPUCopiedBytes << " Bytes)" + << " CPU to GPU: " << numCPUToGPUCopies + << " (" << numCPUToGPUCopiedBytes << " Bytes)"; + } + RDMA_LOG(2) << "Copying tensor " << key + << " From: " << src_addr << " To: " << dst_addr; +#endif // RDMA_COUNT_COPIES +} +#endif // GOOGLE_CUDA + +#ifdef RDMA_DATA_VALIDATION +static uint64_t Checksum(Device* device, const DeviceContext* device_context, + const Tensor& in) { + uint64 checksum = 0; + if (DataTypeCanUseMemcpy(in.dtype())) { +#if GOOGLE_CUDA + if (in.TotalBytes() == 0) { + return 0; } - }; - return cb; + checksum = (device_context != nullptr) + ? GPUUtil::Checksum(device, device_context, in) + : GPUUtil::Checksum(in); +#endif // GOOGLE_CUDA + } else { + string s = in.SummarizeValue(999999); + checksum = Hash64(s.c_str(), s.size(), 0); + } + return checksum; } -// Send the next tensor from the buffer's job queue. -void RdmaTensorBuffer::SendNextItem() { - // get the key - string key_with_step_id = ""; - { - mutex_lock lock{mu_}; - if (!queue_.empty()) { - key_with_step_id = queue_.front(); - queue_.pop(); +static void ValidateChecksum(uint64_t expected, uint64_t actual, + const Tensor& in, uint32_t request_index, + const std::string& key, const std::string& msg) { + RDMA_LOG(2) << "Request #" << request_index << ": " << key + << ": Checksum: " << std::hex << " Expected = 0x" << expected + << ". Actual = 0x" << actual << "."; + + if (expected != actual) { + // Checksum failed. There is one case where this is allowed - if the + // tensor is an AssignAdd of the global step. Since the data-validation + // always postpones the Tensor response in order to send a checksum message, + // it is possible that the global-step was updated while the response was + // still in queue. + if ((in.TotalBytes() == 8) && (in.dtype() == DT_INT64)) { + int64_t prev_val = *(int64_t*)DMAHelper::base(&in) - 1; + actual = Hash64((const char*)&prev_val, 8, 0); + } + if (expected != actual) { + LOG(FATAL) << "[" << msg << "]: Checksum validation failed for request #" + << request_index << ": " << key << std::hex << " " + << DataTypeString(in.dtype()) << " " + << in.shape().DebugString() << " (0x" << in.TotalBytes() + << " bytes): " + << " Expected 0x" << expected << ". Got 0x" << actual << "."; } } +} +#endif // RDMA_DATA_VALIDATION + +#if GOOGLE_CUDA +// Sync the 'done' operation on the GPU stream, but without all the data +// copying. +static void StreamGPUOp(Device* gpu_device, + const DeviceContext* device_context, + StatusCallback done) { + Tensor dummy1, dummy2; + GPUUtil::CopyGPUTensorToCPU( + gpu_device, device_context, &dummy1, &dummy2, done); +} +#endif // GOOGLE_CUDA + +RdmaTensorResponse* RdmaChannel::AddTensorResponse(const RdmaMessage& rm) { + mutex_lock lock{mu_}; + auto it = + responses_table_.emplace(rm.request_index_, RdmaTensorResponse(this, rm)); + CHECK(it.second) << "Response with the ID " << rm.request_index_ + << " already exists."; + return &it.first->second; +} + +RdmaTensorResponse* RdmaChannel::UpdateTensorResponse(const RdmaMessage& rm) { + mutex_lock lock{mu_}; + auto it = responses_table_.find(rm.request_index_); + CHECK(it != responses_table_.end()) << "No response found."; + RdmaTensorResponse* response = &it->second; + response->Update(rm); + return response; +} - // send the tensor if a key is acquired. - if (key_with_step_id != "") { - VLOG(2) << "try to send tensor: " << key_with_step_id; - string key; - int64 step_id; - VerbsUtil::GetKeyAndStepId(key_with_step_id, key, step_id); - CHECK(key.compare(name_) == 0); - Rendezvous::ParsedKey parsed; - Rendezvous::ParseKey(key, &parsed); - Rendezvous::DoneCallback cb = - getRecvTensorCallback(key_with_step_id, key, step_id, parsed); - channel_->adapter_->worker_env_->rendezvous_mgr->RecvLocalAsync(step_id, - parsed, cb); +void RdmaChannel::RemoveTensorResponse(uint32_t request_index) { + mutex_lock lock{mu_}; + responses_table_.erase(request_index); +} + +void RdmaTensorResponse::Start() { + Rendezvous::ParsedKey parsed; + Status s = Rendezvous::ParseKey(rm_.name_, &parsed); + if (!s.ok()) { + SendErrorStatus(s); + return; } + + channel_->adapter_->worker_env_->rendezvous_mgr->RecvLocalAsync( + rm_.step_id_, parsed, + [this, parsed](const Status& status, const Rendezvous::Args& send_args, + const Rendezvous::Args& recv_args, const Tensor& in, + bool is_dead) { + CHECK(status.ok()) << "RecvLocalAsync was not ok." + << " error message: " << status.error_message(); + RecvHandler(parsed, send_args, recv_args, in, is_dead); + }); } -void RdmaTensorBuffer::ReSendNextItem() { - // get the key - string key_with_step_id = ""; - { - mutex_lock lock{mu_}; - if (!requeue.empty()) { - key_with_step_id = requeue.front(); - requeue.pop(); - } +void RdmaTensorResponse::Resume() { SendContent(*tensor_, *proto_, is_dead_); } + +// Helper for RecvTensor. Validates "key" and returns the source +// device in "*src_dev". +Status RdmaTensorResponse::PrepareRecvTensor( + const Rendezvous::ParsedKey& parsed, Device** src_dev) { + // Figures out which device the tensor is hosted on. + string local_name = DeviceNameUtils::LocalName(parsed.src_device); + TF_RETURN_IF_ERROR(channel_->adapter_->worker_env_->device_mgr->LookupDevice( + local_name, src_dev)); + + // Does the device have the right incarnation number we expect? + if ((*src_dev)->attributes().incarnation() != parsed.src_incarnation) { + return errors::Aborted( + "RecvTensor expects a different device incarnation: ", + parsed.src_incarnation, " vs. ", (*src_dev)->attributes().incarnation(), + ". Your worker job was probably restarted. Check your " + "worker job for the reason why it was restarted."); } - // send the tensor if a key is acquired. - if (key_with_step_id != "") { - VLOG(2) << "try to send tensor: " << key_with_step_id; - string key; - int64 step_id; - VerbsUtil::GetKeyAndStepId(key_with_step_id, key, step_id); - CHECK(key.compare(name_) == 0); - Rendezvous::ParsedKey parsed; - Rendezvous::ParseKey(key, &parsed); - Rendezvous::DoneCallback cb = - getRecvTensorCallback(key_with_step_id, key, step_id, parsed); - ReItem* item; - { - mutex_lock lock{mu_}; - Itable it = retable.find(key_with_step_id); - CHECK(it != retable.end()) << "Could not find dup-recv context"; - item = it->second; - retable.erase(it); + return Status::OK(); +} + +void RdmaTensorResponse::RecvHandler(Rendezvous::ParsedKey parsed, + const Rendezvous::Args& send_args, + const Rendezvous::Args& recv_args, + const Tensor& in, bool is_dead) { + Status s = PrepareRecvTensor(parsed, &src_dev_); + if (!s.ok()) { + SendErrorStatus(s); + return; + } + + meta_data_changed_ = TensorMetaDataChanged(in, is_dead); +#ifdef RDMA_DATA_VALIDATION + // Always send a meta data message with the source checksum + meta_data_changed_ = rm_.type_ == RDMA_MESSAGE_TENSOR_REQUEST; + checksum_ = Checksum(src_dev_, send_args.device_context, in); +#endif + bool can_memcpy = DataTypeCanUseMemcpy(in.dtype()); + // string tensor needs to be serialized + Tensor copy; + TensorProto proto; + const bool on_host = send_args.alloc_attrs.on_host(); + if (src_dev_->tensorflow_gpu_device_info() && !on_host) { +#if GOOGLE_CUDA + DeviceContext* send_dev_context = send_args.device_context; + CHECK(send_dev_context) + << "send dev name: " << src_dev_->name() + << " gpu_info: " << src_dev_->tensorflow_gpu_device_info(); + + if (can_memcpy) { + // If the tensor is located on a GDR compatible GPU, there is no need to + // copy it. We can send directly from the source, just need to make sure + // we are in sync with the GPU stream. + // If the tensor's meta-data changed however, we will need to clone it, + // so anyway we'll have to copy it from GPU to CPU first. If at some + // point in time Clone() is changed to only save a shallow copy, we can + // skip the copy here as well. + if ((in.TotalBytes() > 0) && !meta_data_changed_ && + (RdmaMemoryMgr::Singleton().FindMemoryRegion( + (void*)DMAHelper::base(&in), in.TotalBytes()) != nullptr)) { + StreamGPUOp(src_dev_, send_dev_context, + [this, in, proto, is_dead](const Status& s) { + Send(in, proto, is_dead, s); + }); + return; + } + + // The tensor must be copied from GPU to CPU, because either: + // 1. The tensor is located on a non GDR compatible GPU. + // 2. The tensor's meta-data has changed. + Allocator* alloc = ProcessState::singleton()->GetCUDAHostAllocator(0); + copy = Tensor(alloc, in.dtype(), in.shape()); + CountCopies(rm_.name_, (void*)DMAHelper::base(&in), + (void*)DMAHelper::base(©), in.TotalBytes(), true); + GPUUtil::CopyGPUTensorToCPU( + src_dev_, send_dev_context, &in, ©, + [this, copy, proto, is_dead](const Status& s) { + Send(copy, proto, is_dead, s); + }); + } else { + GPUUtil::SetProtoFromGPU( + in, src_dev_, send_args.device_context, &proto, is_dead, + [this, in, proto, is_dead](const Status& s) mutable { + Send(in, proto, is_dead, s); + }); + } +#else + SendErrorStatus(errors::Internal("No GPU device in process")); +#endif // GOOGLE_CUDA + } else { + // tensor is in CPU memory. + if (!can_memcpy) { + in.AsProtoTensorContent(&proto); } - cb(Status::OK(), item->send_args, item->recv_args, item->in, item->is_dead); - delete (item); + Send(in, proto, is_dead, Status::OK()); + } +} + +void RdmaTensorResponse::Send(const Tensor& in, const TensorProto& proto, + bool is_dead, const Status& status) { + if (!status.ok()) { + SendErrorStatus(status); + return; + } + bool can_memcpy = DataTypeCanUseMemcpy(in.dtype()); + bool proto_size_changed = (!can_memcpy) && + (proto.ByteSize() != rm_.tensor_bytes_); + if (meta_data_changed_ || proto_size_changed) { + Clone(in, proto, is_dead); + SendMetaData(in, proto, is_dead); + } else { + SendContent(in, proto, is_dead); + } +} + +bool RdmaTensorResponse::TensorMetaDataChanged(const Tensor& in, bool is_dead) { + return (rm_.data_type_ != in.dtype()) || (rm_.tensor_shape_ != in.shape()) || + (rm_.is_dead_ != is_dead); +} + +void RdmaTensorResponse::Clone(const Tensor& in, const TensorProto& proto, + bool is_dead) { + // Clone the data to be sent later. For simplicity, we clone the tensor's + // data even if it is already a copy. Performance is less of a concern here + // since the meta-data hardly ever changes. The reason we create a copy, is + // that some tensors share their buffer between different step-ids, so the + // tensor content may change before re-request was completed. + bool can_memcpy = DataTypeCanUseMemcpy(in.dtype()); + if (can_memcpy && (in.TotalBytes() > 0)) { + AllocatorAttributes host_alloc_attrs; + host_alloc_attrs.set_nic_compatible(true); + host_alloc_attrs.set_on_host(true); + Allocator* allocator = src_dev_->GetAllocator(host_alloc_attrs); + tensor_ = new Tensor(allocator, in.dtype(), in.shape()); + memcpy(DMAHelper::base(tensor_), DMAHelper::base(&in), in.TotalBytes()); + } else { + tensor_ = new Tensor(in.dtype(), in.shape()); } + if (!can_memcpy) { + proto_ = new TensorProto(proto); + } + is_dead_ = is_dead; } -void RdmaTensorBuffer::PostCopyOperations( - bool can_memcpy, size_t buffer_size, size_t tensor_bytes, const string& key, - const Tensor& in, int64 step_id, bool is_dead, - const string& key_with_step_id, const Tensor* copy, - const TensorProto* proto, const StringPiece* copy_buf, - const Rendezvous::Args& send_args, const Rendezvous::Args& recv_args) { - // prepare message +void RdmaTensorResponse::SendMetaData(const Tensor& in, + const TensorProto& proto, bool is_dead) { + RDMA_LOG(2) << "Request #" << rm_.request_index_ + << ": Meta data changed: " << rm_.name_; + bool can_memcpy = DataTypeCanUseMemcpy(in.dtype()); + size_t tensor_bytes = (can_memcpy) ? in.TotalBytes() : proto.ByteSize(); + + // Send meta-data update: RdmaMessage rm; - rm.name_size_ = key.size(); - rm.name_ = key; + rm.type_ = RDMA_MESSAGE_META_DATA_UPDATE; + rm.name_size_ = rm_.name_.size(); + rm.name_ = rm_.name_; rm.tensor_shape_ = in.shape(); rm.data_type_ = in.dtype(); - rm.step_id_ = step_id; + rm.step_id_ = rm_.step_id_; rm.is_dead_ = is_dead; rm.tensor_bytes_ = tensor_bytes; - rm.buffer_size_ = buffer_size; - mu_.lock(); - if (local_status_ == none || (buffer_size > size_ && local_status_ == idle && - remote_status_ == idle)) { - if ((local_status_ != none) && (buffer_size > size_)) { - VLOG(2) << "Extend RDMA buffer from " << size_ << " to " << buffer_size; - } - CreateCPUBuffer(buffer_size, false); - // Need to be received again, put into the re-recv queue and the table - requeue.push(key_with_step_id); - ReItem* item = new ReItem(send_args, recv_args, in, is_dead); - retable.insert(std::pair<string, ReItem*>(key_with_step_id, item)); - mu_.unlock(); - // no longer used: put back the key since it is not sent; - // ask the remote to create the same buffer - rm.type_ = RDMA_MESSAGE_BUFFER_REQUEST; - rm.remote_addr_ = reinterpret_cast<uint64_t>(buffer_); - rm.rkey_ = self_->rkey; - string message = RdmaMessage::CreateMessage(rm); - channel_->tx_message_buffer_->EnqueueItem(message); - channel_->tx_message_buffer_->SendNextItem(); - } else if ((local_status_ == idle) && (remote_status_ == idle)) { - // both buffers are ready, send the tensor - local_status_ = busy; - remote_status_ = busy; - // local/remote_status_ won't be set back to idle - // unitl Write() is successful - mu_.unlock(); - if (!((buffer_size == size_ && rm.data_type_ != DT_STRING) || - (buffer_size <= size_ && rm.data_type_ == DT_STRING))) { - VLOG(2) << "Tensor and buffer size do not agree," - << " buffer_size = " << size_ - << " requested tensor size = " << buffer_size << in.DebugString(); - } - uint32_t imm_data = LookupBufferIndex(key); - rm.type_ = RDMA_MESSAGE_TENSOR_WRITE; - string message = RdmaMessage::CreateMessage(rm); - memcpy(buffer_, message.data(), message.size()); - if (!is_dead) { - // copy the tensor buffer content - void* output = static_cast<void*>(static_cast<char*>(buffer_) + - RdmaMessage::kTensorBufferStartIndex); - CHECK(tensor_bytes + RdmaMessage::kTensorBufferStartIndex <= size_); - if (can_memcpy) { - CHECK(copy != NULL) << "callback missing pointer to copy tensor"; - CHECK(copy_buf != NULL) << "callback missing pointer to copy buffer"; - CHECK(copy_buf->size() == tensor_bytes) - << "unexpected tensor size: " << copy_buf->size() - << " != " << tensor_bytes; - memcpy(output, copy_buf->data(), tensor_bytes); - } else { - CHECK(proto != NULL) << "callback missing pointer to proto tensor"; - proto->SerializeToArray(output, tensor_bytes); + rm.request_index_ = rm_.request_index_; +#ifdef RDMA_DATA_VALIDATION + rm.checksum_ = checksum_; +#endif + RDMA_LOG(1) << "Step 0x" << std::hex << rm.step_id_ << std::dec + << ": Sending RDMA_MESSAGE_META_DATA_UPDATE #" + << rm.request_index_ << ": " << rm.name_ + << " (shape = " << rm.tensor_shape_.DebugString() << "." + << " data-type = " << DataTypeString(rm.data_type_) << "." + << " is-dead = " << rm.is_dead_ << ")"; + + string message = RdmaMessage::CreateMessage(rm); + channel_->tx_message_buffer_->EnqueueItem(message); + channel_->tx_message_buffer_->SendNextItem(); +} + +void RdmaTensorResponse::SendContent(const Tensor& in, const TensorProto& proto, + bool is_dead) { + bool can_memcpy = DataTypeCanUseMemcpy(in.dtype()); + size_t tensor_bytes = (can_memcpy) ? in.TotalBytes() : proto.ByteSize(); + uint32_t imm_data = rm_.request_index_; + if (!is_dead) { + if (can_memcpy) { + src_buffer_ = const_cast<TensorBuffer*>(DMAHelper::buffer(&in)); + if (src_buffer_ != nullptr) { + src_buffer_->Ref(); // Keep buffer alive until write is complete + src_addr_ = src_buffer_->data(); + mr_ = RdmaMemoryMgr::Singleton().FindMemoryRegion(src_addr_, + tensor_bytes); } } else { - buffer_size = RdmaMessage::kMessageTotalBytes; + RDMA_LOG(2) << "Encoding proto: " << rm_.name_ + << " (Size: " << tensor_bytes << ") " << in.DebugString(); + src_addr_ = malloc(tensor_bytes); + mr_ = ibv_reg_mr(channel_->adapter_->pd_, src_addr_, tensor_bytes, + IBV_ACCESS_LOCAL_WRITE | IBV_ACCESS_REMOTE_WRITE); + proto.SerializeToArray(src_addr_, tensor_bytes); } - Write(imm_data, buffer_size); } else { - // Need to be received again, put into the re-recv queue and the table - requeue.push(key_with_step_id); - ReItem* item = new ReItem(send_args, recv_args, in, is_dead); - retable.insert(std::pair<string, ReItem*>(key_with_step_id, item)); - mu_.unlock(); + tensor_bytes = 0; } + + uint32_t lkey = (mr_ == nullptr) ? 0 : mr_->lkey; + RDMA_LOG(1) << "Step 0x" << std::hex << rm_.step_id_ << std::dec + << ": Sending tensor content #" << rm_.request_index_ << " from " + << std::hex << src_addr_ << " (0x" << lkey << ")" + << " to " << rm_.remote_addr_ << " (0x" << rm_.rkey_ + << "): " << rm_.name_ << " (size: 0x" << std::hex << tensor_bytes + << ")"; + + RdmaMessageBuffer::Write(channel_, imm_data, tensor_bytes, + (uint64_t)src_addr_, lkey, rm_.remote_addr_, + rm_.rkey_, RDMA_WRITE_ID_TENSOR_WRITE, this); +} + +void RdmaTensorResponse::SendErrorStatus(const Status& status) { + RdmaMessage rm; + rm.type_ = RDMA_MESSAGE_ERROR_STATUS; + rm.name_size_ = rm_.name_.size(); + rm.name_ = rm_.name_; + rm.step_id_ = rm_.step_id_; + rm.request_index_ = rm_.request_index_; + rm.status_ = status; + LOG(ERROR) << "Step 0x" << std::hex << rm.step_id_ << std::dec + << ": Sending RDMA_MESSAGE_ERROR_STATUS #" + << rm.request_index_ << ": " << rm.name_ + << ". Status: " << status.ToString(); + + string message = RdmaMessage::CreateMessage(rm); + channel_->tx_message_buffer_->EnqueueItem(message); + channel_->tx_message_buffer_->SendNextItem(); + + // Destroy the response. + Destroy(); +} + +void RdmaTensorResponse::Destroy() { + if (src_buffer_ != nullptr) { + src_buffer_->Unref(); + } + if (tensor_ != nullptr) { + delete tensor_; + } + if (proto_ != nullptr) { + ibv_dereg_mr(mr_); + free(src_addr_); + delete proto_; + } + // Remove response from the pending list: + channel_->RemoveTensorResponse(rm_.request_index_); } // Create a RdmaMessage according to the pre-defined format @@ -1276,43 +1273,46 @@ void RdmaTensorBuffer::PostCopyOperations( // message in string format string RdmaMessage::CreateMessage(const RdmaMessage& rm) { // Rdma Message format - // type|name_size|name|step_id|buffer_size|remote_addr|rkey|is_dead|... - // 1B| 2B | 512| 8B | 8B | 8B | 4B | 1B |... - // ...|data_type|tensor_shape|tensor_bytes|tensor_buffer - // ...| XB | XB | 8B |... + // type|name_size|name|step_id|request_index|remote_addr|rkey|is_dead|... + // 1B| 2B | 512| 8B | 8B | 8B | 4B | 1B |... + // ...|data_type|tensor_shape|tensor_bytes|error_status | + // ...| XB | XB | 8B |size - 4B, proto - XB | // - // ACK: type|13|"rx_ack_buffer" - // TENSOR_REQUEST: type|name_size|tensor_name|step_id - // TENSOR_WRITE: type|name_size|tensor_name|step_id|...|is_dead - // |data_type|tensor_shape|tensor_bytes - // BUFFER_IDLE: type|name_size|buffer_name - // BUFFER_REQUEST: - // type|name_size|buffer_name|...|buffer_size|remote_addr|rkey| - // BUFFER_RESPONSE: - // type|name_size|buffer_name|...|buffer_size|remote_addr|rkey| - char message[kMessageTotalBytes]; + // ACK: Imm-type: ACK + // TENSOR_REQUEST: Imm-type: MESSAGE + // Fields: type, request_index, name, step_id, remote_addr, + // rkey, is_dead, data_type, tensor_shape, tensor_bytes + // META_DATA_UPDATE: Imm-type: MESSAGE + // Fields: type, request_index, is_dead, data_type, + // tensor_shape, tensor_bytes + // TENSOR_RE_REQUST: Imm-type: MESSAGE + // Fields: type, request_index, name, step_id, remote_addr, + // rkey, is_dead, data_type, tensor_shape, tensor_bytes + // ERROR_STATUS: Imm-type: MESSAGE + // Fields: type, request_index, name, step_id, error_status + // Tensor content: Imm-type: request_index + size_t message_size = kMessageTotalBytes; + char message[kMessageTotalBytes + kErrorStatusMaxSize]; // type message[kTypeStartIndex] = static_cast<char>(rm.type_) & 0xff; - // size of name - memcpy(&message[kNameSizeStartIndex], &rm.name_size_, sizeof(rm.name_size_)); - // name - memcpy(&message[kNameStartIndex], rm.name_.data(), rm.name_.size()); - // buffer_size, remote_addr, rkey - if ((rm.type_ == RDMA_MESSAGE_BUFFER_REQUEST) || - (rm.type_ == RDMA_MESSAGE_BUFFER_RESPONSE)) { - memcpy(&message[kBufferSizeStartIndex], &rm.buffer_size_, - sizeof(rm.buffer_size_)); + // request index + memcpy(&message[kRequestIndexStartIndex], &rm.request_index_, + sizeof(rm.request_index_)); + // name, step_id, remote_addr, rkey + if ((rm.type_ == RDMA_MESSAGE_TENSOR_REQUEST) || + (rm.type_ == RDMA_MESSAGE_TENSOR_RE_REQUEST)) { + memcpy(&message[kNameSizeStartIndex], &rm.name_size_, + sizeof(rm.name_size_)); + memcpy(&message[kNameStartIndex], rm.name_.data(), rm.name_.size()); memcpy(&message[kRemoteAddrStartIndex], &rm.remote_addr_, sizeof(rm.remote_addr_)); memcpy(&message[kRkeyStartIndex], &rm.rkey_, sizeof(rm.rkey_)); - } - // step_id - if ((rm.type_ == RDMA_MESSAGE_TENSOR_WRITE) || - (rm.type_ == RDMA_MESSAGE_TENSOR_REQUEST)) { memcpy(&message[kStepIdStartIndex], &rm.step_id_, sizeof(rm.step_id_)); } // is_dead, data_type, tensor_shape, tensor_bytes - if (rm.type_ == RDMA_MESSAGE_TENSOR_WRITE) { + if ((rm.type_ == RDMA_MESSAGE_TENSOR_REQUEST) || + (rm.type_ == RDMA_MESSAGE_META_DATA_UPDATE) || + (rm.type_ == RDMA_MESSAGE_TENSOR_RE_REQUEST)) { memcpy(&message[kIsDeadStartIndex], &rm.is_dead_, sizeof(rm.is_dead_)); memcpy(&message[kDataTypeStartIndex], &rm.data_type_, @@ -1322,7 +1322,31 @@ string RdmaMessage::CreateMessage(const RdmaMessage& rm) { memcpy(&message[kTensorBytesStartIndex], &rm.tensor_bytes_, sizeof(rm.tensor_bytes_)); } - return string(message, kMessageTotalBytes); + // checksum +#ifdef RDMA_DATA_VALIDATION + memcpy(&message[kChecksumStartIndex], &rm.checksum_, sizeof(rm.checksum_)); +#endif + // error status + if (rm.type_ == RDMA_MESSAGE_ERROR_STATUS) { + ::grpc::Status gs = ToGrpcStatus(rm.status_); + ErrorStatusProto gsProto; + gsProto.set_error_code(gs.error_code()); + gsProto.set_error_message(gs.error_message()); + gsProto.set_error_details(gs.error_details()); + uint32_t gsProtoSize = gsProto.ByteSize(); + if (gsProtoSize + 4 > kErrorStatusMaxSize) { + LOG(ERROR) << "Error status (" << gsProtoSize + 4 << " bytes) " + << "is too big to fit in RDMA message (" + << kErrorStatusMaxSize << " bytes). Truncated."; + gsProtoSize = kErrorStatusMaxSize - 4; + } + uint32_t* proto_size = (uint32_t*)&message[kErrorStatusStartIndex]; + *proto_size = gsProtoSize; + gsProto.SerializeToArray(&message[kErrorStatusStartIndex + 4], + gsProtoSize); + message_size += gsProtoSize + 4; + } + return string(message, message_size); } // Parse a RdmaMessage according to the pre-defined format @@ -1335,26 +1359,24 @@ void RdmaMessage::ParseMessage(RdmaMessage& rm, void* buffer) { char* message = static_cast<char*>(buffer); // type rm.type_ = static_cast<RdmaMessageType>(message[kTypeStartIndex]); - // name_size_ - memcpy(&rm.name_size_, &message[kNameSizeStartIndex], sizeof(rm.name_size_)); - // name - rm.name_ = string(&message[kNameStartIndex], rm.name_size_); - // buffer_size, remote_addr, rkey - if ((rm.type_ == RDMA_MESSAGE_BUFFER_REQUEST) || - (rm.type_ == RDMA_MESSAGE_BUFFER_RESPONSE)) { - memcpy(&rm.buffer_size_, &message[kBufferSizeStartIndex], - sizeof(rm.buffer_size_)); + // request index + memcpy(&rm.request_index_, &message[kRequestIndexStartIndex], + sizeof(rm.request_index_)); + // name, step_id, remote_addr, rkey + if ((rm.type_ == RDMA_MESSAGE_TENSOR_REQUEST) || + (rm.type_ == RDMA_MESSAGE_TENSOR_RE_REQUEST)) { + memcpy(&rm.name_size_, &message[kNameSizeStartIndex], + sizeof(rm.name_size_)); + rm.name_ = string(&message[kNameStartIndex], rm.name_size_); memcpy(&rm.remote_addr_, &message[kRemoteAddrStartIndex], sizeof(rm.remote_addr_)); memcpy(&rm.rkey_, &message[kRkeyStartIndex], sizeof(rm.rkey_)); - } - // step_id - if ((rm.type_ == RDMA_MESSAGE_TENSOR_WRITE) || - (rm.type_ == RDMA_MESSAGE_TENSOR_REQUEST)) { memcpy(&rm.step_id_, &message[kStepIdStartIndex], sizeof(rm.step_id_)); } // data_type, tensor_bytes, tensor_shape, is_dead - if (rm.type_ == RDMA_MESSAGE_TENSOR_WRITE) { + if ((rm.type_ == RDMA_MESSAGE_TENSOR_REQUEST) || + (rm.type_ == RDMA_MESSAGE_META_DATA_UPDATE) || + (rm.type_ == RDMA_MESSAGE_TENSOR_RE_REQUEST)) { memcpy(&rm.is_dead_, &message[kIsDeadStartIndex], sizeof(rm.is_dead_)); memcpy(&rm.data_type_, &message[kDataTypeStartIndex], sizeof(rm.data_type_)); @@ -1363,6 +1385,294 @@ void RdmaMessage::ParseMessage(RdmaMessage& rm, void* buffer) { memcpy(&rm.tensor_bytes_, &message[kTensorBytesStartIndex], sizeof(rm.tensor_bytes_)); } + // checksum +#ifdef RDMA_DATA_VALIDATION + memcpy(&rm.checksum_, &message[kChecksumStartIndex], sizeof(rm.checksum_)); +#endif + // error status + if (rm.type_ == RDMA_MESSAGE_ERROR_STATUS) { + ErrorStatusProto gsProto; + uint32_t gsProtoSize = *(uint32_t*)&message[kErrorStatusStartIndex]; + CHECK(ParseProtoUnlimited( + &gsProto, &message[kErrorStatusStartIndex + 4], gsProtoSize)) + << "Failed to parse error status proto from message. Aborting."; + ::grpc::Status gs((::grpc::StatusCode)gsProto.error_code(), + gsProto.error_message(), gsProto.error_details()); + rm.status_ = FromGrpcStatus(gs); + } +} + +//***************************************************************************** +// RdmaMemoryMgr +//***************************************************************************** + +ibv_mr* RdmaMemoryMgr::FindMemoryRegion(void* addr, size_t length) { + mutex_lock l(mrs_mu_); + auto iter = std::upper_bound(mrs_.begin(), mrs_.end(), addr, &Comparator); + if (iter == std::end(mrs_) || iter->get()->addr > addr) { + return nullptr; + } else { + return iter->get(); + } +} + +void RdmaMemoryMgr::InsertMemoryRegion(void* addr, size_t length, + const std::string& allocator_name) { + if (length == 0) return; + ibv_mr* mr = ibv_reg_mr(pd_, addr, length, + IBV_ACCESS_LOCAL_WRITE | IBV_ACCESS_REMOTE_WRITE); + RDMA_LOG(1) << "Insert memory region 0x" << std::hex << mr->rkey << ". [" + << addr << "-" << (void*)((uint64_t)addr + length - 1) << "]" + << " SIZE: 0x" << length << " (" << allocator_name << ")."; + if (mr != nullptr) { + mutex_lock l(mrs_mu_); + auto iter = std::upper_bound(mrs_.begin(), mrs_.end(), addr, &Comparator); + mrs_.insert(iter, {mr, &MRDeleter}); + } else { + LOG(WARNING) << "Cannot register memory region"; + } +} + +void RdmaMemoryMgr::EvictMemoryRegion(void* addr, size_t length) { + if (length == 0) return; + mutex_lock l(mrs_mu_); + auto iter = std::upper_bound(mrs_.begin(), mrs_.end(), addr, &Comparator); + if (iter != std::end(mrs_) && iter->get()->addr == addr) { + mrs_.erase(iter); + RDMA_LOG(1) << "Evict memory region 0x" << std::hex << iter->get()->rkey; + + } else { + LOG(WARNING) << "Failed to de-register memory region"; + } +} + +const TensorMetaData* RdmaMemoryMgr::GetTensorMetaData( + const std::string& tensor_name) { + mutex_lock l(tensor_meta_data_mu_); + auto it = tensors_meta_data_.find(tensor_name); + if (it == tensors_meta_data_.end()) { + return nullptr; + } + return &it->second; +} + +const TensorMetaData* RdmaMemoryMgr::SetTensorMetaData( + const std::string& tensor_name, DataType dtype, const TensorShape& shape, + bool is_dead, size_t proto_size) { + mutex_lock l(tensor_meta_data_mu_); + TensorMetaData& meta_data = tensors_meta_data_[tensor_name]; + meta_data.data_type_ = dtype; + meta_data.tensor_shape_ = shape; + meta_data.proto_size_ = proto_size; + meta_data.is_dead_ = is_dead; + return &meta_data; +} + +//***************************************************************************** +// RdmaTensorRequest +//***************************************************************************** + +RdmaTensorRequest::RdmaTensorRequest( + uint32_t index, const string& key, int64 step_id, RdmaChannel* channel, + Device* dst_dev, const Rendezvous::Args recv_args, + const RdmaTensorRequest::RecvDoneCallback& done) + : index_(index), + key_(key), + step_id_(step_id), + channel_(channel), + dst_dev_(dst_dev), + recv_args_(recv_args), + meta_data_(RdmaMemoryMgr::Singleton().GetTensorMetaData(key)), + result_tensor_(nullptr), + proxy_tensor_(nullptr), + rdma_addr_(nullptr), + mr_(nullptr), + done_(done) {} + +RdmaTensorRequest::~RdmaTensorRequest() { DeallocateTensors(); } + +void RdmaTensorRequest::Done(const Status& s) { + Tensor val = std::move(*result_tensor_); + +#ifdef RDMA_DATA_VALIDATION + // Validate checksum + // Unfortunately we can't always do a Checksum directly on the result tensor. + // If the result tensor is on GPU, then we need to copy it back to CPU. If + // we happen to be in the midst of a proxy callback, then the copying will + // get stuck. + uint64_t checksum = (proxy_tensor_ != nullptr) + ? Checksum(nullptr, nullptr, *proxy_tensor_) + : Checksum(dst_dev_, recv_args_.device_context, val); + ValidateChecksum(checksum_, checksum, val, index_, key_, "RDMA"); +#endif + + Rendezvous::Args recv_args = std::move(recv_args_); + bool is_dead = (meta_data_ == nullptr) ? false : meta_data_->is_dead_; + RecvDoneCallback done = done_; + DeallocateTensors(); + channel_->RemoveTensorRequest(index_); + done(s, Rendezvous::Args(), recv_args, val, is_dead); +} + +void RdmaTensorRequest::DeallocateTensors() { + if (result_tensor_ != nullptr) { + delete result_tensor_; + result_tensor_ = nullptr; + } + if (proxy_tensor_ != nullptr) { + delete proxy_tensor_; + proxy_tensor_ = nullptr; + } +} + +bool RdmaTensorRequest::AllocateTensors() { + result_tensor_ = + new Tensor(dst_dev_->GetAllocator(recv_args_.alloc_attrs), + meta_data_->data_type_, meta_data_->tensor_shape_); + + size_t tensor_size = result_tensor_->TotalBytes(); + bool can_memcpy = DataTypeCanUseMemcpy(result_tensor_->dtype()); + if (can_memcpy) { + if (tensor_size == 0) { + return true; + } + rdma_addr_ = DMAHelper::base(result_tensor_); + mr_ = RdmaMemoryMgr::Singleton().FindMemoryRegion(rdma_addr_, tensor_size); +#if GOOGLE_CUDA + if (mr_ == nullptr) { + // Can't RDMA directly to result. Use a proxy. + proxy_tensor_ = + new Tensor(ProcessState::singleton()->GetCUDAHostAllocator(0), + result_tensor_->dtype(), result_tensor_->shape()); + rdma_addr_ = DMAHelper::base(proxy_tensor_); + mr_ = + RdmaMemoryMgr::Singleton().FindMemoryRegion(rdma_addr_, tensor_size); + } +#endif + } else { + uint32_t proto_size = meta_data_->proto_size_; + rdma_addr_ = malloc(proto_size); + mr_ = ibv_reg_mr(RdmaMemoryMgr::Singleton().pd_, rdma_addr_, proto_size, + IBV_ACCESS_LOCAL_WRITE | IBV_ACCESS_REMOTE_WRITE); + } + CHECK(mr_ != nullptr) << " No memory region found for address " << rdma_addr_ + << ": " << key_; + return true; +} + +void RdmaTensorRequest::AllocateTensorsAsync(StatusCallback done) { + AllocateTensors(); + bool on_host = recv_args_.alloc_attrs.on_host(); + if (dst_dev_->tensorflow_gpu_device_info() && !on_host && + (proxy_tensor_ == nullptr)) { +#if GOOGLE_CUDA + // We need to sync the memory allocation on the GPU: + StreamGPUOp(dst_dev_, recv_args_.device_context, done); +#endif + } else { + done(Status::OK()); + } +} + +void RdmaTensorRequest::Send(RdmaMessageType message_type) { + RdmaMessageBuffer* rb = channel_->tx_message_buffer_; + RdmaMessage rm; + rm.type_ = message_type; + rm.request_index_ = index_; + rm.name_size_ = key_.size(); + rm.name_ = key_; + rm.step_id_ = step_id_; + rm.remote_addr_ = (uint64_t)rdma_addr_; + if (meta_data_ != nullptr) { + rm.data_type_ = meta_data_->data_type_; + rm.tensor_shape_ = meta_data_->tensor_shape_; + rm.is_dead_ = meta_data_->is_dead_; + rm.tensor_bytes_ = meta_data_->proto_size_; + } else { + rm.data_type_ = DT_INVALID; + } + rm.rkey_ = (mr_ == nullptr) ? 0 : mr_->rkey; + + RDMA_LOG(1) << "Step 0x" << std::hex << rm.step_id_ << std::dec + << ": Sending " << MessageTypeToString(message_type) + << " #" << index_ << ": " + << rm.name_ << " on " << rdma_addr_ + << " (rkey: 0x" << std::hex << rm.rkey_ << ")"; + + string message = RdmaMessage::CreateMessage(rm); + rb->EnqueueItem(message); + rb->SendNextItem(); +} + +void RdmaTensorRequest::RecvTensorMetaData(DataType dtype, TensorShape shape, + bool is_dead, size_t proto_size) { + meta_data_ = RdmaMemoryMgr::Singleton().SetTensorMetaData( + key_, dtype, shape, is_dead, proto_size); + + DeallocateTensors(); + AllocateTensorsAsync([this](const Status& s) { + Send(RDMA_MESSAGE_TENSOR_RE_REQUEST); + }); +} + +void RdmaTensorRequest::RecvTensorContent() { + bool can_memcpy = DataTypeCanUseMemcpy(meta_data_->data_type_); + size_t message_size = + can_memcpy ? result_tensor_->TotalBytes() : meta_data_->proto_size_; + RDMA_LOG(1) << "Step 0x" << std::hex << step_id_ << std::dec + << ": Received tensor content #" << index_ << ": " + << key_ << " (Size: 0x" << std::hex << message_size << ")"; + + Tensor val; + +#if GOOGLE_CUDA + if (proxy_tensor_ != nullptr) { + CountCopies(key_, (void*)DMAHelper::base(proxy_tensor_), + (void*)DMAHelper::base(result_tensor_), + result_tensor_->TotalBytes(), false); + GPUUtil::CopyCPUTensorToGPU(proxy_tensor_, recv_args_.device_context, + dst_dev_, result_tensor_, + [this](const Status& s) { + CHECK(s.ok()) << "copy tensor to gpu sync"; + Done(s); + }); + return; + } +#endif + + if (can_memcpy) { + Done(Status::OK()); + } else { + RDMA_LOG(2) << "Decoding proto: " << key_ + << " (Size: " << meta_data_->proto_size_ << ")"; + TensorProto proto; + CHECK(ParseProtoUnlimited(&proto, rdma_addr_, meta_data_->proto_size_)) + << "fail to parse proto from array"; + ibv_dereg_mr(mr_); + free(rdma_addr_); + Status s = dst_dev_->MakeTensorFromProto(proto, recv_args_.alloc_attrs, + result_tensor_); + Done(s); + } +} + +void RdmaTensorRequest::RecvErrorStatus(const Status& status) { + if (result_tensor_ == nullptr) { + result_tensor_ = new Tensor(); + } + LOG(ERROR) << "Received RDMA_MESSAGE_ERROR_STATUS: " << status.ToString(); + Done(status); +} + +void RdmaTensorRequest::Start() { + meta_data_ = RdmaMemoryMgr::Singleton().GetTensorMetaData(key_); + if (meta_data_ != nullptr) { + AllocateTensorsAsync([this](const Status& s) { + Send(RDMA_MESSAGE_TENSOR_REQUEST); + }); + } else { + Send(RDMA_MESSAGE_TENSOR_REQUEST); + } } } // end namespace tensorflow diff --git a/tensorflow/contrib/verbs/rdma.h b/tensorflow/contrib/verbs/rdma.h index fea2327d77..68b3d59f56 100644 --- a/tensorflow/contrib/verbs/rdma.h +++ b/tensorflow/contrib/verbs/rdma.h @@ -27,6 +27,7 @@ limitations under the License. #include <unordered_map> #include <vector> +#include "tensorflow/contrib/verbs/verbs_util.h" #include "tensorflow/core/distributed_runtime/worker_env.h" #include "tensorflow/core/framework/rendezvous.h" #include "tensorflow/core/framework/tensor.h" @@ -43,6 +44,11 @@ namespace tensorflow { #define SL_DEFAULT 0 #define TRAFFIC_CLASS 0 +#define RDMA_LOG_0 LOG(INFO) +#define RDMA_LOG_1 VLOG(1) +#define RDMA_LOG_2 VLOG(2) +#define RDMA_LOG(LEVEL) RDMA_LOG_##LEVEL + struct RdmaParams { uint8_t port_num; uint8_t sgid_index; @@ -76,29 +82,303 @@ enum Location { local, remote }; -enum BufferType { - ACK, - MESSAGE, - TENSOR -}; + enum RdmaMessageType { - RDMA_MESSAGE_ACK, - RDMA_MESSAGE_BUFFER_IDLE, - RDMA_MESSAGE_BUFFER_REQUEST, - RDMA_MESSAGE_BUFFER_RESPONSE, + RDMA_MESSAGE_META_DATA_UPDATE, + RDMA_MESSAGE_TENSOR_RE_REQUEST, RDMA_MESSAGE_TENSOR_REQUEST, - RDMA_MESSAGE_TENSOR_WRITE + RDMA_MESSAGE_ERROR_STATUS, +}; + +struct RdmaMessage { + RdmaMessageType type_; + uint16_t name_size_; + string name_; + int64 step_id_; + uint64_t request_index_; + union { + uint64_t remote_addr_; +#ifdef RDMA_DATA_VALIDATION + uint64_t checksum_; +#endif + }; + uint32_t rkey_; + bool is_dead_; + DataType data_type_; + TensorShape tensor_shape_; + size_t tensor_bytes_; + + // For error status: + Status status_; + + // type|name_size|name|step_id|request_index|remote_addr/checksum|rkey|... + // 1B| 2B | 512| 8B | 8B | 8B | 4B |... + // ...|is_dead|data_type|tensor_shape|tensor_bytes|error_status | + // ...| 1B | XB | XB | 8B |size - 4B, proto - XB | + static const size_t kNameCapacity = 512; + static const size_t kTypeStartIndex = 0; + static const size_t kNameSizeStartIndex = kTypeStartIndex + sizeof(type_); + static const size_t kNameStartIndex = + kNameSizeStartIndex + sizeof(name_size_); + static const size_t kStepIdStartIndex = kNameStartIndex + kNameCapacity; + static const size_t kRequestIndexStartIndex = + kStepIdStartIndex + sizeof(step_id_); + static const size_t kRemoteAddrStartIndex = + kRequestIndexStartIndex + sizeof(request_index_); + static const size_t kChecksumStartIndex = kRemoteAddrStartIndex; + static const size_t kRkeyStartIndex = + kRemoteAddrStartIndex + sizeof(remote_addr_); + static const size_t kIsDeadStartIndex = kRkeyStartIndex + sizeof(rkey_); + static const size_t kDataTypeStartIndex = + kIsDeadStartIndex + sizeof(is_dead_); + static const size_t kTensorShapeStartIndex = + kDataTypeStartIndex + sizeof(data_type_); + static const size_t kTensorBytesStartIndex = + kTensorShapeStartIndex + sizeof(TensorShape); + static const size_t kErrorStatusStartIndex = + kTensorBytesStartIndex + sizeof(tensor_bytes_); + static const size_t kErrorStatusMaxSize = 4096; + + static const size_t kMessageTotalBytes = kErrorStatusStartIndex; + static const size_t kRdmaMessageBufferSize = + kMessageTotalBytes + kErrorStatusMaxSize; + static string CreateMessage(const RdmaMessage& rm); + static void ParseMessage(RdmaMessage& rm, void* buffer); +}; + +// Immediate types for RDMA write +enum RdmaImmDataType { + RDMA_IMM_MAX_REQUEST_ID = 0xFFFFFFFD, + RDMA_IMM_DATA_ACK = 0xFFFFFFFE, + RDMA_IMM_DATA_MESSAGE = 0xFFFFFFFF +}; + +// Write types for RDMA write-complete events +enum RdmaWriteIDType { + RDMA_WRITE_ID_ACK, + RDMA_WRITE_ID_MESSAGE, + RDMA_WRITE_ID_TENSOR_WRITE +}; + +// Context for RDMA write-complete events +class RdmaWriteID { + public: + RdmaWriteID(RdmaWriteIDType write_type, void* write_context) + : write_type(write_type), write_context(write_context) {} + + RdmaWriteIDType write_type; + void* write_context; +}; + +// Tensor meta-data +class TensorMetaData { + public: + TensorShape tensor_shape_; + DataType data_type_; + size_t proto_size_; + bool is_dead_; + + std::ostream& print(std::ostream& out) const { + out << "Dtype = " << DataTypeString(data_type_) + << ", Shape = " << tensor_shape_.DebugString() << ", Proto size = 0x" + << std::hex << proto_size_ << ", Is dead = " << is_dead_; + return out; + } +}; + +inline std::ostream& operator<<(std::ostream& out, + const TensorMetaData& meta_data) { + return meta_data.print(out); +} + +class RdmaChannel; + +void MRDeleter(ibv_mr* mr); +using MemoryRegionPtr = std::unique_ptr<ibv_mr, decltype(&MRDeleter)>; + +// RdmaMemoryMgr +// Manages the local meta-data cache, and the registered RDMA memory regions. +class RdmaMemoryMgr { + public: + static RdmaMemoryMgr& Singleton() { + static RdmaMemoryMgr instance; + return instance; + } + + // Memory regions + ibv_mr* FindMemoryRegion(void* addr, size_t length); + void InsertMemoryRegion(void* addr, size_t length, + const std::string& allocator_name); + void EvictMemoryRegion(void* addr, size_t length); + + // Tensor meta-data cache + const TensorMetaData* GetTensorMetaData(const std::string& tensor_name); + const TensorMetaData* SetTensorMetaData(const std::string& tensor_name, + DataType dtype, + const TensorShape& shape, + bool is_dead, size_t proto_size); + + struct ibv_pd* pd_; + + protected: + RdmaMemoryMgr() : pd_(nullptr) {} + + static bool Comparator(const void* ptr, const MemoryRegionPtr& other) { + return ptr < reinterpret_cast<char*>(other->addr) + other->length; + } + + private: + mutex tensor_meta_data_mu_; + std::unordered_map<std::string, TensorMetaData> tensors_meta_data_; + + // Managed memory regions + mutex mrs_mu_; + std::vector<MemoryRegionPtr> mrs_ GUARDED_BY(mrs_mu_); }; -class RdmaBuffer; + +// RdmaTensorRequest +// Represents a single tensor request. +class RdmaTensorRequest { + public: + typedef Rendezvous::DoneCallback RecvDoneCallback; + + // Creates a tensor request identified by index. + RdmaTensorRequest(uint32_t index, const string& key, int64 step_id, + RdmaChannel* channel, Device* dst_dev, + const Rendezvous::Args recv_args, + const RecvDoneCallback& done); + ~RdmaTensorRequest(); + + // Request unique index. + uint32_t index() { return index_; } + + // Start the tensor request sequence. + // + // 1. Allocate the result tensor (and proxy tensor if required). + // 2. Send RDMA_MESSAGE_TENSOR_REQUEST to the remote side. + void Start(); + + // Receive tensor meta-data. + // + // 1. Update the local meta-data cache. + // 2. Reallocate the result tensor (and proxy tensor if required). + // 3. Re-send the request to the remote side. + void RecvTensorMetaData(DataType dtype, TensorShape shape, bool is_dead, + size_t proto_size); + + // Receive tensor content (RDMA write was completed). + // + // Decode proto if required and/or move to GPU if the content was not + // written to it directly (GPU direct is not avaliable). Afterwards, + // invoke Done(). + void RecvTensorContent(); + + // Receive error status (in case of a remote error). + // Invoke Done() with the status code. + void RecvErrorStatus(const Status& status); + +#ifdef RDMA_DATA_VALIDATION + // Receive tensor checksum + // + // For validation: Get and store the Tensor's expected checksum for the + // current request. Compare the result Tensor's checksum with the stored + // checksum right before invoking Done(). + void RecvTensorChecksum(uint64_t checksum) { checksum_ = checksum; } +#endif + + private: + void Done(const Status& s); + void Send(RdmaMessageType message_type); + bool AllocateTensors(); + void AllocateTensorsAsync(StatusCallback done); + void DeallocateTensors(); + + uint32_t index_; + string key_; + int64 step_id_; + RdmaChannel* channel_; + Device* dst_dev_; + Rendezvous::Args recv_args_; + const TensorMetaData* meta_data_; + Tensor* result_tensor_; + Tensor* proxy_tensor_; + void* rdma_addr_; + ibv_mr* mr_; + RecvDoneCallback done_; +#ifdef RDMA_DATA_VALIDATION + uint64_t checksum_; +#endif +}; + +// RdmaTensorResponse +// Represents a single tensor response. +class RdmaTensorResponse { + public: + // Creates a response for request message. + RdmaTensorResponse(RdmaChannel* channel, const RdmaMessage& rm) + : channel_(channel), rm_(rm) {} + + void Update(const RdmaMessage& rm) { rm_ = rm; } + + // Start the tensor response sequence. + // + // 1. Find the tensor in the local tag-match table and invoke RecvHandler. + // (Using RecvLocalAsync()). + // 2. Compare the tensor's meta-data to the meta-data in the message (taken + // from the requester's local cache). + // If meta-data changed: + // a. Clone the tensor to be sent later. + // b. Send a meta-data update message and wait for re-request. + // Else: + // a. Send the tensor's content (using direct RDMA write). + void Start(); + + // Resume the response sequence, after a re-request. + // + // 1. Send the tensor's content that was cloned earlier. + void Resume(); + + // Destroy the response's resources and remove it from the pending list. + void Destroy(); + + private: + void RecvHandler(Rendezvous::ParsedKey parsed, + const Rendezvous::Args& send_args, + const Rendezvous::Args& recv_args, const Tensor& in, + bool is_dead); + void Clone(const Tensor& in, const TensorProto& proto, bool is_dead); + void Send(const Tensor& in, const TensorProto& proto, bool is_dead, + const Status& status); + bool TensorMetaDataChanged(const Tensor& in, bool is_dead); + Status PrepareRecvTensor(const Rendezvous::ParsedKey& parsed, + Device** src_dev); + void SendMetaData(const Tensor& in, const TensorProto& proto, bool is_dead); + void SendContent(const Tensor& in, const TensorProto& proto, bool is_dead); + void SendErrorStatus(const Status& status); + + RdmaChannel* channel_; + RdmaMessage rm_; // The request message + Device* src_dev_ = nullptr; + TensorBuffer* src_buffer_ = nullptr; + void* src_addr_ = nullptr; + ibv_mr* mr_ = nullptr; + uint64_t checksum_ = 0; + bool meta_data_changed_ = false; + + // Re-item: + TensorProto* proto_ = nullptr; + Tensor* tensor_ = nullptr; + bool is_dead_ = false; +}; + +class RdmaMessageBuffer; // Class that represents the Rdma Adapter. // Responsible for creation of the completion queue, and handling // of work completions. class RdmaAdapter { friend class RdmaChannel; - friend class RdmaBuffer; - friend class RdmaAckBuffer; friend class RdmaMessageBuffer; - friend class RdmaTensorBuffer; + friend class RdmaTensorResponse; friend class RdmaMgr; friend class RdmaRemoteRendezvous; @@ -133,10 +413,10 @@ class RdmaAdapter { // Responsible for connecting queue pairs. class RdmaChannel { friend class RdmaAdapter; - friend class RdmaBuffer; - friend class RdmaAckBuffer; friend class RdmaMessageBuffer; friend class RdmaTensorBuffer; + friend class RdmaTensorRequest; + friend class RdmaTensorResponse; friend class RdmaMgr; friend class RdmaRemoteRendezvous; @@ -146,22 +426,28 @@ class RdmaChannel { ~RdmaChannel(); inline const RdmaAddress& self() { return self_; } RdmaAddress address() const; - inline const std::vector<RdmaBuffer*>& message_buffers() const { + inline const std::vector<RdmaMessageBuffer*>& message_buffers() const { return message_buffers_; } void Connect(const RdmaAddress& remoteAddr); void Connect(); void Recv(); - RdmaBuffer* FindBuffer(const uint32_t index); - RdmaBuffer* FindBuffer(const string& name); - RdmaBuffer* FindOrCreateBuffer(const string& name, - BufferType buffer_type = TENSOR); - uint32_t LookupBufferIndex(const string& buffer_name); void SetRemoteAddress(const RdmaAddress& ra, bool override); - void InsertRecvCallback(const string& key, std::function<void()> recv_done); - void RemoveRecvCallback(const string& key); - void RunRecvCallback(const string& key); - static const int kNumMessageBuffers = 4; + + // Requests: + RdmaTensorRequest* InsertTensorRequest( + const string& key, int64 step_id, Device* dst_dev, + const Rendezvous::Args recv_args, + const RdmaTensorRequest::RecvDoneCallback& done); + void RemoveTensorRequest(uint32_t request_index); + RdmaTensorRequest* GetTensorRequest(uint32_t request_index); + + // Responses: + RdmaTensorResponse* AddTensorResponse(const RdmaMessage& rm); + RdmaTensorResponse* UpdateTensorResponse(const RdmaMessage& rm); + void RemoveTensorResponse(uint32_t request_index); + + static const int kNumMessageBuffers = 2; static const int kPingRecvWrid = 0; private: @@ -179,36 +465,31 @@ class RdmaChannel { string remote_name_; ibv_qp* qp_; mutex mu_; - bool connected_ GUARDED_BY(bt_mu_) = false; - RdmaAddress remote_ GUARDED_BY(bt_mu_); - bool remote_set_ GUARDED_BY(bt_mu_) = false; + bool connected_ GUARDED_BY(mu_) = false; + RdmaAddress remote_ GUARDED_BY(mu_); + bool remote_set_ GUARDED_BY(mu_) = false; mutex ct_mu_; - typedef std::unordered_map<string, std::function<void()> > CallbackTable; - CallbackTable callback_table_ GUARDED_BY(ct_mu_); - mutex bt_mu_; - typedef std::unordered_map<unsigned int, RdmaBuffer*> BufferTable; - BufferTable buffer_table_ GUARDED_BY(bt_mu_); - typedef std::unordered_map<uint32_t, string> BufferIndexNameTable; - BufferIndexNameTable buffer_index_name_table_ GUARDED_BY(bt_mu_); - typedef std::unordered_map<string, uint32_t> BufferNameIndexTable; - BufferNameIndexTable buffer_name_index_table_ GUARDED_BY(bt_mu_); - RdmaBuffer* tx_message_buffer_; - RdmaBuffer* rx_message_buffer_; - RdmaBuffer* tx_ack_buffer_; - RdmaBuffer* rx_ack_buffer_; - std::vector<RdmaBuffer*> message_buffers_; + typedef std::unordered_map<uint32_t, RdmaTensorRequest> RequestTable; + RequestTable request_table_ GUARDED_BY(ct_mu_); + uint32_t request_serial_ GUARDED_BY(ct_mu_); + mutex responses_mu_; + typedef std::unordered_map<uint32_t, RdmaTensorResponse> ResponsesTable; + ResponsesTable responses_table_ GUARDED_BY(responses_mu_); + RdmaMessageBuffer* tx_message_buffer_; + RdmaMessageBuffer* rx_message_buffer_; + std::vector<RdmaMessageBuffer*> message_buffers_; }; -// Class that represents a buffer for Rdma writes and reads. -class RdmaBuffer { +// Class that represents a buffer for Rdma message sending. +class RdmaMessageBuffer { friend class RdmaChannel; friend class RdmaAdapter; friend class RdmaMgr; friend class RdmaRemoteRendezvous; public: - explicit RdmaBuffer(RdmaChannel* channel, string name); - virtual ~RdmaBuffer(); + explicit RdmaMessageBuffer(RdmaChannel* channel, string name); + ~RdmaMessageBuffer(); inline void* buffer() const { return buffer_; } inline ibv_mr* self() const { return self_; } @@ -223,13 +504,15 @@ class RdmaBuffer { } void FreeBuffer(); void EnqueueItem(string Item); - virtual void SendNextItem() {}; + void SendNextItem(); void CreateCPUBuffer(size_t size, bool lock = true); void SetRemoteMR(RemoteMR rmi, bool override); - uint32_t LookupBufferIndex(const string& buffer_name) { - return const_cast<RdmaChannel*>(channel_)->LookupBufferIndex(buffer_name); - } void Write(uint32_t imm_data, size_t buffer_size); + static void Write(const RdmaChannel* channel, uint32_t imm_data, + size_t buffer_size, uint64_t src_addr, uint32_t lkey, + uint64_t remote_addr, uint32_t rkey, + RdmaWriteIDType write_type, void* write_context); + static void SendAck(const RdmaChannel* channel); protected: const RdmaChannel* channel_; @@ -245,125 +528,6 @@ class RdmaBuffer { BufferStatus remote_status_ GUARDED_BY(mu_) = none; }; -class RdmaAckBuffer : public RdmaBuffer { - public: - explicit RdmaAckBuffer(RdmaChannel* channel, string name); - virtual ~RdmaAckBuffer() override {} - void SendNextItem() override; -}; - -class RdmaMessageBuffer : public RdmaBuffer { - friend class RdmaChannel; - friend class RdmaAapater; - - public: - explicit RdmaMessageBuffer(RdmaChannel* channel, string name); - virtual ~RdmaMessageBuffer() override {} - void SendNextItem() override; -}; - -class RdmaTensorBuffer : public RdmaBuffer { - public: - explicit RdmaTensorBuffer(RdmaChannel* channel, string name); - virtual ~RdmaTensorBuffer() override; - void SendNextItem() override; - void PostCopyOperations(bool can_memcpy, size_t buffer_size, - size_t tensor_bytes, const string& key, - const Tensor& in, int64 step_id, bool is_dead, - const string& key_with_step_id, const Tensor* copy, - const TensorProto* proto, const StringPiece* copy_buf, - const Rendezvous::Args& send_args, - const Rendezvous::Args& recv_args); - - void ReSendNextItem(); - - private: - Rendezvous::DoneCallback getRecvTensorCallback( - const string& key_with_step_id, const string& key, int64 step_id, - const Rendezvous::ParsedKey& parsed); - - struct ReItem { - Rendezvous::Args send_args; - Rendezvous::Args recv_args; - Tensor in; - bool is_dead; - - ReItem(const Rendezvous::Args& send_args_, - const Rendezvous::Args& recv_args_, const Tensor& in_, bool is_dead_) - : send_args(send_args_), - recv_args(recv_args_), - in(in_), - is_dead(is_dead_) { - if (send_args.device_context) { - send_args.device_context->Ref(); - } - if (recv_args.device_context) { - recv_args.device_context->Ref(); - } - } - - ~ReItem() { - if (send_args.device_context) { - send_args.device_context->Unref(); - } - if (recv_args.device_context) { - recv_args.device_context->Unref(); - } - } - }; - typedef std::map<string, ReItem*> Table; - typedef Table::iterator Itable; - - std::queue<string> requeue GUARDED_BY(mu_); - Table retable GUARDED_BY(mu_); -}; - -struct RdmaMessage { - RdmaMessageType type_; - uint16_t name_size_; - string name_; - int64 step_id_; - uint64_t buffer_size_; - uint64_t remote_addr_; - uint32_t rkey_; - bool is_dead_; - DataType data_type_; - TensorShape tensor_shape_; - size_t tensor_bytes_; - - // type|name_size|name|step_id|buffer_size|remote_addr|rkey|is_dead|... - // 1B| 2B | 512| 8B | 8B | 8B | 4B | 1B |... - // ...|data_type|tensor_shape|tensor_bytes|tensor_buffer - // ...| XB | XB | 8B |... - // - static const size_t kNameCapacity = 512; - static const size_t kTypeStartIndex = 0; - static const size_t kNameSizeStartIndex = kTypeStartIndex + sizeof(type_); - static const size_t kNameStartIndex = - kNameSizeStartIndex + sizeof(name_size_); - static const size_t kStepIdStartIndex = kNameStartIndex + kNameCapacity; - static const size_t kBufferSizeStartIndex = - kStepIdStartIndex + sizeof(step_id_); - static const size_t kRemoteAddrStartIndex = - kBufferSizeStartIndex + sizeof(buffer_size_); - static const size_t kRkeyStartIndex = - kRemoteAddrStartIndex + sizeof(remote_addr_); - static const size_t kIsDeadStartIndex = kRkeyStartIndex + sizeof(rkey_); - static const size_t kDataTypeStartIndex = - kIsDeadStartIndex + sizeof(is_dead_); - static const size_t kTensorShapeStartIndex = - kDataTypeStartIndex + sizeof(data_type_); - static const size_t kTensorBytesStartIndex = - kTensorShapeStartIndex + sizeof(TensorShape); - static const size_t kTensorBufferStartIndex = - kTensorBytesStartIndex + sizeof(tensor_bytes_); - static const size_t kMessageTotalBytes = kTensorBufferStartIndex; - static const size_t kRdmaMessageBufferSize = kMessageTotalBytes; - static const size_t kRdmaAckBufferSize = kMessageTotalBytes; - static string CreateMessage(const RdmaMessage& rm); - static void ParseMessage(RdmaMessage& rm, void* buffer); -}; - } // namespace tensorflow #endif // TENSORFLOW_USE_VERBS diff --git a/tensorflow/contrib/verbs/rdma_mgr.cc b/tensorflow/contrib/verbs/rdma_mgr.cc index 9cb307bcfa..f3644af0b4 100644 --- a/tensorflow/contrib/verbs/rdma_mgr.cc +++ b/tensorflow/contrib/verbs/rdma_mgr.cc @@ -16,11 +16,16 @@ limitations under the License. #ifdef TENSORFLOW_USE_VERBS #include "tensorflow/contrib/verbs/rdma_mgr.h" +#include <fstream> #include <vector> #include "tensorflow/contrib/verbs/grpc_verbs_client.h" #include "tensorflow/contrib/verbs/verbs_service.pb.h" +#include "tensorflow/core/common_runtime/bfc_allocator.h" +#include "tensorflow/core/common_runtime/gpu/gpu_util.h" +#include "tensorflow/core/common_runtime/gpu/process_state.h" #include "tensorflow/core/distributed_runtime/rpc/grpc_worker_cache.h" #include "tensorflow/core/distributed_runtime/session_mgr.h" +#include "tensorflow/core/framework/allocator_registry.h" #include "tensorflow/core/lib/core/status.h" namespace tensorflow { @@ -53,7 +58,7 @@ RdmaMgr::RdmaMgr(const WorkerEnv* const worker_env, void RdmaMgr::SetupChannels() { for (const auto& p : channel_table_) { string worker_name = p.first; - LOG(INFO) << "connecting to remote node " << worker_name; + RDMA_LOG(2) << "Connecting to remote node " << worker_name; RdmaChannel* rc = p.second; GetRemoteAddressRequest req; GetRemoteAddressResponse resp; @@ -78,39 +83,49 @@ void RdmaMgr::SetupChannels() { mr->set_rkey(rc->message_buffers_[i]->self_->rkey); } // synchronous call - Status s = client->GetRemoteAddress(&req, &resp); - // save obtained remote addresses - // connect to the remote channel - if (s.ok()) { - CHECK(worker_name.compare(resp.host_name()) == 0); - RdmaAddress ra; - ra.lid = resp.channel().lid(); - ra.qpn = resp.channel().qpn(); - ra.psn = resp.channel().psn(); - ra.snp = resp.channel().snp(); - ra.iid = resp.channel().iid(); - rc->SetRemoteAddress(ra, false); - rc->Connect(); - int i = 0; - int idx[] = {1, 0, 3, 2}; - for (const auto& mr : resp.mr()) { - // the connections are crossed, i.e. - // local tx_message_buffer <---> remote rx_message_buffer_ - // local rx_message_buffer <---> remote tx_message_buffer_ - // local tx_ack_buffer <---> remote rx_ack_buffer_ - // local rx_ack_buffer <---> remote tx_ack_buffer_ - // hence idx[] = {1, 0, 3, 2}. - RdmaBuffer* rb = rc->message_buffers_[idx[i]]; - RemoteMR rmr; - rmr.remote_addr = mr.remote_addr(); - rmr.rkey = mr.rkey(); - rb->SetRemoteMR(rmr, false); - i++; + Status s; + int attempts = 0; + static const int max_num_attempts = 5; + do { + s = client->GetRemoteAddress(&req, &resp); + // save obtained remote addresses + // connect to the remote channel + if (s.ok()) { + CHECK(worker_name.compare(resp.host_name()) == 0); + RdmaAddress ra; + ra.lid = resp.channel().lid(); + ra.qpn = resp.channel().qpn(); + ra.psn = resp.channel().psn(); + ra.snp = resp.channel().snp(); + ra.iid = resp.channel().iid(); + rc->SetRemoteAddress(ra, false); + rc->Connect(); + int i = 0; + int idx[] = {1, 0}; + for (const auto& mr : resp.mr()) { + // the connections are crossed, i.e. + // local tx_message_buffer <---> remote rx_message_buffer_ + // local rx_message_buffer <---> remote tx_message_buffer_ + // hence idx[] = {1, 0}. + RdmaMessageBuffer* rb = rc->message_buffers_[idx[i]]; + RemoteMR rmr; + rmr.remote_addr = mr.remote_addr(); + rmr.rkey = mr.rkey(); + rb->SetRemoteMR(rmr, false); + i++; + } + CHECK(i == RdmaChannel::kNumMessageBuffers); + } else { + LOG(ERROR) << "Connecting to " << worker_name + << ": Got " << s.error_message() << ". Retrying (" + << (attempts + 1) << "/" << max_num_attempts << ")..." ; + if (++attempts == max_num_attempts) { + break; + } + worker_env_->env->SleepForMicroseconds(2000000); } - CHECK(i == RdmaChannel::kNumMessageBuffers); - } else { - LOG(ERROR) << s.error_message(); - } + } while (!s.ok()); + RDMA_LOG(0) << "Connected to remote node " << worker_name; delete client; } } @@ -183,6 +198,138 @@ RdmaChannel* RdmaMgr::FindChannel(const string& name) { return iter->second; } +bool IsGDRAvailable() { +#if defined(__APPLE__) + return false; +#elif defined(PLATFORM_WINDOWS) + return false; +#else + std::ifstream ifs("/proc/modules"); + string line; + while (std::getline(ifs, line)) { + auto sep = line.find(' '); + CHECK_NE(sep, std::string::npos); + if (line.substr(0, sep) == "nv_peer_mem") { + return true; + } + } + return false; +#endif +} + +int TryToReadNumaNode(ibv_device* device) { +#if defined(__APPLE__) + LOG(INFO) << "OS X does not support NUMA - returning NUMA node 0"; + return 0; +#elif defined(PLATFORM_WINDOWS) + // Windows support for NUMA is not currently implemented. Return node 0. + return 0; +#else + VLOG(2) << "Trying to read NUMA node for device: " << device->name; + static const int kUnknownNumaNode = -1; + + auto filename = string(device->ibdev_path) + "/device/numa_node"; + + std::ifstream ifs(filename.c_str()); + string content; + CHECK(std::getline(ifs, content)); + + int32 value; + if (strings::safe_strto32(content, &value)) { + if (value < 0) { + LOG(INFO) << "Successful NUMA node read from SysFS had negative value (" + << value << "), but there must be at least one NUMA node" + ", so returning NUMA node zero"; + return 0; + } + LOG(INFO) << "NUMA node for device: " << device->name << " is " << value; + return value; + } + return kUnknownNumaNode; +#endif +} + +void MRDeleter(ibv_mr* mr) { + if (mr) { + ibv_dereg_mr(mr); + } +} + +// TODO(byronyi): remove this class duplicated from the one in +// common/runtime/gpu/pool_allocator.h when it is available in common_runtime +class BasicCPUAllocator : public SubAllocator { + public: + ~BasicCPUAllocator() override {} + + void* Alloc(size_t alignment, size_t num_bytes) override { + return port::AlignedMalloc(num_bytes, alignment); + } + void Free(void* ptr, size_t) override { port::AlignedFree(ptr); } +}; + +// TODO(byronyi): remove this class and its registration when the default +// cpu_allocator() returns visitable allocator +class BFCRdmaAllocator : public BFCAllocator { + public: + BFCRdmaAllocator() + : BFCAllocator(new BasicCPUAllocator(), 1LL << 36, true, "cpu_rdma_bfc") { + } +}; + +REGISTER_MEM_ALLOCATOR("BFCRdmaAllocator", 101, BFCRdmaAllocator); + +void RdmaMgr::InitAllocators() { + RdmaMemoryMgr::Singleton().pd_ = rdma_adapter_->pd_; + + Allocator* allocators[] = { +#if GOOGLE_CUDA + ProcessState::singleton()->GetCUDAHostAllocator(0), + ProcessState::singleton()->GetCPUAllocator(0), +#endif // GOOGLE_CUDA + cpu_allocator(), + }; + + using namespace std::placeholders; + + std::set<Allocator*> instrumented_; + + // Host memory allocators + for (Allocator* allocator : allocators) { + VisitableAllocator::Visitor alloc_visitor = + std::bind(&RdmaMemoryMgr::InsertMemoryRegion, + &RdmaMemoryMgr::Singleton(), _1, _2, allocator->Name()); + VisitableAllocator::Visitor free_visitor = std::bind( + &RdmaMemoryMgr::EvictMemoryRegion, &RdmaMemoryMgr::Singleton(), _1, _2); + + auto* visitable_allocator = dynamic_cast<VisitableAllocator*>(allocator); + CHECK(visitable_allocator) << "is not visitable for instrumentation" + << allocator->Name(); + // Make sure we don't instrument the same allocator twice + if (instrumented_.find(allocator) == std::end(instrumented_)) { + visitable_allocator->AddAllocVisitor(alloc_visitor); + visitable_allocator->AddFreeVisitor(free_visitor); + instrumented_.insert(allocator); + LOG(INFO) << "Instrumenting CPU allocator " << allocator->Name(); + } + } + +#if GOOGLE_CUDA + if (IsGDRAvailable()) { + // Note we don't free allocated GPU memory so there is no free visitor + int32_t bus_id = TryToReadNumaNode(rdma_adapter_->context_->device) + 1; + + char buf[8]; + sprintf(buf, "gpu"); + VisitableAllocator::Visitor cuda_alloc_visitor = + std::bind(&RdmaMemoryMgr::InsertMemoryRegion, + &RdmaMemoryMgr::Singleton(), _1, _2, std::string(buf)); + + ProcessState::singleton()->AddGPUAllocVisitor(bus_id, cuda_alloc_visitor); + LOG(INFO) << "Instrumenting GPU allocator with bus_id " << bus_id; + } +#endif // GOOGLE_CUDA +} + } // end namespace tensorflow #endif diff --git a/tensorflow/contrib/verbs/rdma_mgr.h b/tensorflow/contrib/verbs/rdma_mgr.h index e711e60478..29227a9544 100644 --- a/tensorflow/contrib/verbs/rdma_mgr.h +++ b/tensorflow/contrib/verbs/rdma_mgr.h @@ -38,6 +38,7 @@ class RdmaMgr { RdmaChannel* FindChannel(const string& key); void SetupChannels(); bool ConnectivityCheck(); + void InitAllocators(); const string& local_worker() { return local_worker_; } private: diff --git a/tensorflow/contrib/verbs/rdma_rendezvous_mgr.cc b/tensorflow/contrib/verbs/rdma_rendezvous_mgr.cc index 74f6681af3..ad3dce1784 100644 --- a/tensorflow/contrib/verbs/rdma_rendezvous_mgr.cc +++ b/tensorflow/contrib/verbs/rdma_rendezvous_mgr.cc @@ -21,10 +21,6 @@ limitations under the License. #include "tensorflow/core/common_runtime/device.h" #include "tensorflow/core/common_runtime/device_mgr.h" #include "tensorflow/core/common_runtime/dma_helper.h" -#if GOOGLE_CUDA -#include "tensorflow/core/common_runtime/gpu/gpu_util.h" -#include "tensorflow/core/common_runtime/gpu/process_state.h" -#endif // GOOGLE_CUDA #include "tensorflow/core/lib/core/errors.h" #include "tensorflow/core/lib/strings/numbers.h" #include "tensorflow/core/lib/strings/str_util.h" @@ -36,11 +32,6 @@ class RdmaRemoteRendezvous : public BaseRemoteRendezvous { RdmaRemoteRendezvous(const WorkerEnv* env, int64 step_id, RdmaMgr* rdma_mgr) : BaseRemoteRendezvous(env, step_id), rdma_mgr_(rdma_mgr) {} - void RecvPostCopyOps(const string& key, const string& key_with_step_id, - const Rendezvous::Args& recv_args, - const DoneCallback& done, const RdmaMessage& rm, - RdmaChannel* rc, Tensor& val, const Status& s); - protected: void RecvFromRemoteAsync(const Rendezvous::ParsedKey& parsed, const Rendezvous::Args& args, @@ -74,101 +65,18 @@ void RdmaRemoteRendezvous::RecvFromRemoteAsync( RdmaChannel* rc = rdma_mgr_->FindChannel(src_name); string key(std::move(parsed.FullKey().ToString())); string key_with_step_id = VerbsUtil::AppendStepidToKey(key, step_id_); - // insert callback - rc->InsertRecvCallback(key_with_step_id, [this, key, key_with_step_id, rc, - recv_args, parsed, done]() { - Status src_s, dst_s, s; - Device* src_dev, *dst_dev; - src_s = env_->device_mgr->LookupDevice("CPU:0", &src_dev); - dst_s = env_->device_mgr->LookupDevice(parsed.dst_device, &dst_dev); - if (!src_s.ok() || !dst_s.ok()) { - s = src_s.ok() ? dst_s : src_s; - LOG(ERROR) << "s is not ok, error code " << s.error_message(); - done(s, Args(), recv_args, Tensor(), true); - return; - } - RdmaBuffer* rb = rc->FindBuffer(key); - RdmaMessage rm; - CHECK(rb->size_ >= RdmaMessage::kMessageTotalBytes); - RdmaMessage::ParseMessage(rm, rb->buffer_); - CHECK(rm.type_ == RDMA_MESSAGE_TENSOR_WRITE); - Tensor val; - if (!rm.is_dead_) { - void* input = static_cast<char*>(rb->buffer_) + - RdmaMessage::kTensorBufferStartIndex; - bool can_memcpy = DataTypeCanUseMemcpy(rm.data_type_); - if (can_memcpy) { - if (dst_dev->tensorflow_gpu_device_info() && - (!recv_args.alloc_attrs.on_host())) { -#if GOOGLE_CUDA - CHECK(recv_args.device_context) - << "send dev name: " << src_dev->name() - << " gpu_info: " << src_dev->tensorflow_gpu_device_info(); - Allocator* alloc = ProcessState::singleton()->GetCUDAHostAllocator(0); - Tensor copy(alloc, rm.data_type_, rm.tensor_shape_); - memcpy(DMAHelper::base(©), input, rm.tensor_bytes_); - - Allocator* dst_alloc = dst_dev->GetAllocator(recv_args.alloc_attrs); - Tensor gpu_copy(dst_alloc, rm.data_type_, rm.tensor_shape_); - - GPUUtil::CopyCPUTensorToGPU( - ©, recv_args.device_context, dst_dev, &gpu_copy, - [this, gpu_copy, key, key_with_step_id, recv_args, done, rm, rc]( - const Status& s) { - CHECK(s.ok()) << "copy tensor to gpu sync"; - Tensor val; - val = std::move(gpu_copy); - RecvPostCopyOps(key, key_with_step_id, recv_args, done, rm, rc, - val, s); - }); -#endif // GOOGLE_CUDA - return; - } else { - AllocatorAttributes host_alloc_attrs; - host_alloc_attrs.set_gpu_compatible(true); - host_alloc_attrs.set_on_host(true); - Allocator* alloc = dst_dev->GetAllocator(host_alloc_attrs); - Tensor copy(alloc, rm.data_type_, rm.tensor_shape_); - memcpy(DMAHelper::base(©), input, rm.tensor_bytes_); - val = std::move(copy); - } - } else { - TensorProto proto; - CHECK(rm.tensor_bytes_ + RdmaMessage::kTensorBufferStartIndex <= - rb->size_); - CHECK(ParseProtoUnlimited(&proto, input, rm.tensor_bytes_)) - << "fail to parse proto from array"; - s = dst_dev->MakeTensorFromProto(proto, recv_args.alloc_attrs, &val); - } - } - RecvPostCopyOps(key, key_with_step_id, recv_args, done, rm, rc, val, s); - }); - // append key to message queue - RdmaBuffer* rb = rc->tx_message_buffer_; - RdmaMessage rm; - rm.type_ = RDMA_MESSAGE_TENSOR_REQUEST; - rm.name_size_ = key.size(); - rm.name_ = key; - rm.step_id_ = step_id_; - string message = RdmaMessage::CreateMessage(rm); - rb->EnqueueItem(message); - rb->SendNextItem(); -} -void RdmaRemoteRendezvous::RecvPostCopyOps( - const string& key, const string& key_with_step_id, - const Rendezvous::Args& recv_args, const DoneCallback& done, - const RdmaMessage& rm, RdmaChannel* rc, Tensor& val, const Status& s) { - rc->RemoveRecvCallback(key_with_step_id); - RdmaMessage br; - br.type_ = RDMA_MESSAGE_BUFFER_IDLE; - br.name_size_ = key.size(); - br.name_ = key; - string message = RdmaMessage::CreateMessage(br); - RdmaBuffer* tb = rc->tx_message_buffer_; - tb->EnqueueItem(message); - tb->SendNextItem(); - done(s, Args(), recv_args, val, rm.is_dead_); + Device* dst_dev; + s = env_->device_mgr->LookupDevice(parsed.dst_device, &dst_dev); + CHECK(s.ok()) << "s is not ok, error code " << s.error_message(); + if (!s.ok()) { + done(s, Args(), recv_args, Tensor(), true); + return; + } + + RdmaTensorRequest* request = + rc->InsertTensorRequest(key, step_id_, dst_dev, recv_args, done); + request->Start(); } RdmaRendezvousMgr::RdmaRendezvousMgr(const WorkerEnv* env) diff --git a/tensorflow/contrib/verbs/verbs_server_lib.cc b/tensorflow/contrib/verbs/verbs_server_lib.cc index a606ef75a4..47ed83f521 100644 --- a/tensorflow/contrib/verbs/verbs_server_lib.cc +++ b/tensorflow/contrib/verbs/verbs_server_lib.cc @@ -104,6 +104,7 @@ Status VerbsServer::Start() { [this] { verbs_service_->HandleRPCsLoop(); })); rdma_mgr_->SetupChannels(); CHECK(rdma_mgr_->ConnectivityCheck()) << "Connectivity check failed!"; + rdma_mgr_->InitAllocators(); verbs_state_ = CONNECTED; } } diff --git a/tensorflow/contrib/verbs/verbs_service.proto b/tensorflow/contrib/verbs/verbs_service.proto index 0df1fed4b9..abdae1d84f 100644 --- a/tensorflow/contrib/verbs/verbs_service.proto +++ b/tensorflow/contrib/verbs/verbs_service.proto @@ -50,6 +50,12 @@ message GetRemoteAddressResponse { repeated MemoryRegion mr = 3; } +message ErrorStatusProto { + int32 error_code = 1; + string error_message = 2; + string error_details = 3; +} + //////////////////////////////////////////////////////////////////////////////// // // VerbsService diff --git a/tensorflow/contrib/verbs/verbs_with_0_copies.png b/tensorflow/contrib/verbs/verbs_with_0_copies.png Binary files differnew file mode 100644 index 0000000000..0641e2fd50 --- /dev/null +++ b/tensorflow/contrib/verbs/verbs_with_0_copies.png diff --git a/tensorflow/contrib/verbs/verbs_with_0_copies.xml b/tensorflow/contrib/verbs/verbs_with_0_copies.xml new file mode 100644 index 0000000000..16130a961b --- /dev/null +++ b/tensorflow/contrib/verbs/verbs_with_0_copies.xml @@ -0,0 +1 @@ +<mxfile userAgent="Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/63.0.3239.84 Safari/537.36" version="7.8.7" editor="www.draw.io" type="device"><diagram name="Page-1" id="74e2e168-ea6b-b213-b513-2b3c1d86103e">7Vxtc9o4EP41zKQfmsGW3/hIgPQ60/RyIZ1rPzHClsFXY1FZEOivP8mW8ZsAB2yXtHQ6jb2SJXl3n0e7K6cdMFhsPhC4nD9gB/kdtetsOmDYUVVFUw32g0u2scRUu7FgRjxHdEoFY+8nEsKk28pzUJjrSDH2qbfMC20cBMimORkkBL/ku7nYz8+6hDNUEoxt6Jel/3oOnQup0u2mDX8hbzYXU1u6aJhC+/uM4FUg5uuowI3+xM0LmIwl+odz6OCXjAiMOmBAMKbx1WIzQD7XbaK2+Ln7Pa27dRMU0CoP6CB+Yg39FUqWHC2MbhNlRK+D+APdDrh7mXsUjZfQ5q0vzPxMNqcLn90p7NL1fH+AfUzYfYAD1ulOzIAIRZu9y1R2L8+cCuEFomTLumx2mo8fEf5kiduX1DhWIptn7GIkQigcYrYbOlUKuxB6kevIkqjI8NkMd463zqnK+LHihrtjL0rfQ9+bBR3QZz185NK0lV3NxM9olHAJg0Q2ppDQm3dJE1tatjUjjqbOS+tfTSKbEskK2lqYcsua+r6PbUgRJwIUhJiEt03OqfI5xyhwbp5Hn8d/P02eRv98GY2f37VhAam2c7PVBVDGTg5Elmtzu1OCv6NMi2FbaOru5iuBVQLp/fgFefwqRhnAalcC4B3jngPghGxrR/ATstfPkTs+IAqHkMKbC/GQ2oD3ZenEsGNah+/ZNeTbLrTnqHkAPiHYLuxlHAjKVCBjg8q0+XsBGahtAlkxjkcryGGRnLjFhM7xDAfQH6XSu7yWMxr9D1G6FcEoXFHMROkInzBein579RjiFbGTEFIsjW3oM5R002IZX+NBbRPkQ+qt89HoWZpTG6LAiCQGBMUgJShc4iBshRvs9SfGDX4/3AZ2s7QbUcAAL7dRGsKvH79wyCPisWd/8hf/6EZv/2PlEeTsffs3B3dT+aX7tv4r0M20RbZfszff+GC3Or/dePRr7u6bmKgaJ2hlTkhS5fo4QTz6iD22lJ0pe728KU2jYKF4UeKp1Eh9QuA2023JO4QH5jELO4RZyECP9E9SttRH4hWkHrPTSTXm00rMd++RkD57Cw7cjjngf1UDLjjggmm4LPJGjN4kwhvMYTBDDmMccF8V53O8mK7C4xh3GHvY1MOMjYbMbzjCLgP3LW/zvePbfDiHS37p+mjT5xWfCKyOuBzaPgxDzz5Ioa5lI9vOIr5bRnxJzVNL1/TKiLckQYBaEfAZXesSVSeyM3kBFCI6tcgL8euUeKE0kNbt3jK/MUxLUSxD10wjP65SjW9OgHjiHm35y5k+Id0FuhflFIbSu1WtHlAVy1KApqs5U2pFU1Z1kcPDgoo70ikeUi5zfkNhyUl4OJh3gbypRUFTUuMUMeTQZoZHTH7Hudbj4aloWHiOE8Unsi0gH7M0wd+gzN+axH3UOqK28ob7Gf++qraK8U6bqpblw00VQqJM75GFyfI0r61yMM/+5MFadgVPwwc+xAvxorz0Jq7SdaITI8qM/e61S3/zqZsh8cvmQjigH9+S28zlRMK2i+2q7tWqKdmrW8rYrKIFtWr74/6M7Zwd1CwZdFcaztDBm4eJ1mqFA1Q4fn0jmU6yn+WQYlZESjtRrVpIdTTzxDi2OJBe9IWaaimleZR6ayOgQj29EfeTlNbOdT+j7H7gstzvcPZjgkaSqtIXEPUlVaC8JdR9rDqIo7Xf6VRVUlqMI2uCN9soQOXnDPcOyh4veJWOF0oDyyLj6PTkY7BmYGMXDs+p+IGu7/NPl45Fxa9S0ZEz1aHkdJf7aeBE77rAayReGoX0tEzjzQfxxePWdoP4JN6UAHyuVIKAyNFlCEeMSEnGUHzEPXY6vVbAXMX2ghkT6Ondc5QevFf3GZ05HnH9aHebe46DgibKBoqp5yzbKxt299Fb1rDFHOAku6pN2ZV/J/EnW9U0jlq115RRZamEYO0iL7o4CiDsHWmldgSu2+1y8ihBdvjQnzyMxuP+h9Ek/1Vcxt7xyCUWnjbgBRdWBwRWnqoVyTeq0uiyjkKgVq65Nmf8h9FzfzLss3+eRuPHvz+PR1cHkDgA0Np0AFm9rXH0XwnggP21Vglg/0lALfYv1s+vFpdY3NDbtHhT2XfNSXUi+LhYxP2ruCF3Qv47M8VRBQO9yvsOJYiyVLLm9773kO+E+VfPLpBulyzPHaSp7sRjXrqJRQFciMaQouWE2V70XGCKJtBxiBB8R9v4in+mPYk/0z6SGZ9w6nUMuTwvNqaGbnTKNUDXVaMa4KVh2MhjWJV86QRkGbZeB4ab/tWihjAsg1m31PvPBbpUPweBfoXtebAFkmCrOdjKvk+8wvYPhO3r9ucrsk9Att7mhpyMcUV2ceqC818+viueVF1xtwd3hqR2XRfu2G36XxzEJ8/p/yMBRv8D</diagram></mxfile>
\ No newline at end of file diff --git a/tensorflow/contrib/verbs/verbs_with_0_copies_phase1_protocol.jpg b/tensorflow/contrib/verbs/verbs_with_0_copies_phase1_protocol.jpg Binary files differnew file mode 100644 index 0000000000..8bc69b889d --- /dev/null +++ b/tensorflow/contrib/verbs/verbs_with_0_copies_phase1_protocol.jpg diff --git a/tensorflow/contrib/verbs/verbs_with_0_copies_phase1_protocol.xml b/tensorflow/contrib/verbs/verbs_with_0_copies_phase1_protocol.xml new file mode 100644 index 0000000000..484e7c78ae --- /dev/null +++ b/tensorflow/contrib/verbs/verbs_with_0_copies_phase1_protocol.xml @@ -0,0 +1 @@ +<mxfile userAgent="Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/62.0.3202.94 Safari/537.36" version="7.8.4" editor="www.draw.io" type="device"><diagram name="Page-1" id="74e2e168-ea6b-b213-b513-2b3c1d86103e">7Vxbc5s4FP41nuk+pMMd8ujYTrYzTZqNk9nmyaOAbNhgRIWc2P31K4G4yzZ1AHtaZzJjOBLS8TnnOzeRDNTRcn2DQejeIgf6A0Vy1gN1PFAUWVMM+sEom4RiGpcJYYE9h0/KCVPvJ+REiVNXngOj0kSCkE+8sEy0URBAm5RoAGP0Xp42R3551xAsYI0wtYFfp/7rOcTlVFmS8oG/obdw+daWzgdegP26wGgV8P0GijqPf5LhJUjX4vMjFzjovUBSJwN1hBEiydVyPYI+k20qtuS56y2jGd8YBqTJA1bywBvwVzDl2PDpo1eO98b4IxsuE+PHijF1ReCaXADfWwQDdUhn+HBO8lF6teCf8SpRCIKUNiUAk09/pUOUqeJogRxvXaa2z01Ke8ECDnpkDCxDehG8ROzjgs4c+j6yAYGPMIgQjkoC64WBKQycT4+Tu+m3h9nD5J+nyfSxax62a6K0m1LaRYlxBpklS3T43fUInIbAZqPv1C9RmkuWPr2Ts6ffIKZMbQWLnEGQujaIlpDgDZ3CH7A4aLlTkw1+/567CCX1EG7BO2RuA3C3tMiWzqFJLzg6xUhNPUbrUH2A9ltia7eQgDEgoHOTa6juNnZjBj2GoE9M1UHc/X4xZq+erq8nDLPT+29308nWTU8LRqrSJ4xkQwCjikCgQ5MBfoswcdECBcCf5NSrssgK4vkPErLh+QxYEURJ+QpfEQpLYmQb7RYi5QutsJ3O4rzSQLqA6TRNLGwMfUC8t/L6H5Kc0pEDivHiON/wU+hQyDzAKERBBLsHDfN8XylM/WG0Cezu9/syj/XyY+VhZjvDPgP7gKGsJRL7LiMUbm7unx7R6F6UXT2dUp7XqSCmEHuUsZ/wqN/45HMnQztq8qQfw8lT0eDN9+LNM1vss85u1x75Xrp75hsdGBq0emhIqvAPhAb+6D3y6M6ZKi+VsipNo6KhhAf+VK6kIcZgU5gWsglR831Us1LL7plvWLvnW9rO+fQi4Ti3sEyGzQKmVguY1x6OyKO3hMyZmCP2W/MyBbeQYDdNy0cuCBbQoX5GvW6KchctX1bRfoQ7NCTZxEPU2YypWTFEdoH6nnO9y/25XuSCkF3Ofbgess5RDFWHX45tH0SRZ5eFNfd8f4R8hOMl0gZPAe9SHe+HodoS5HuKWOIFieoCgaa0D2K/mrwrVewn3NewX1tI1aXPpiwZpiXLlqFrplFeV27mUw6AZWoEfVlFi/5cGhxR9bpx+VmxLlVFtixZ1XSlpDCtqrCmhrB7WbVhbDnEDtSaHTzDqGYKLA8rKzoiGL3CVNUBCmBF+5zEk7exTfUMKf2KuVKPlRt8YOk5TpxpiJxzOfvowherdV+sCcxHaSP/qofCO/T7itqSjihqUYOjq+KKFUD3KBSW7H2VQBfbUqgiAw/jW7bCO6bap5+fkr7cID5BIlSrv8r4iVVThsDAusurVH1/BO2zvOI1VJZwHRx0FVMQdLsposyqBrVmgW57EfWRUGjWFLqlIXdidq/12kVQ6xlDTSKnXU+kAaZk4KZY5P1klXKloNDMA/PI6kJ6VeMtdSVq+8jtdg3UBgcUnRiZoEl1oJEZdSNTj2pku2sMU+2kdEnbSR2ULmrdX7d9FjxK8qLf6ShY0Fo78FSmvyOWETtiubl/aqSHftgaw0h45LGbq/hJWqzte+TEEm3rmHl2muwIYO7KjYDA62ERziF1p7ggdbbiFqEfXpfTUsr2ggUl6PndY5zBXyjbNIgoY3M/jmQuLdth0E2JXlLsZUO9VrP0g9QqakC2olb2FsifrNRqedCrVkXFAQ9vVSc3R3EaYWfpWK5ImphJEuOxBtnx7XB2O5lOhzeTWfntvILCk5VrLvWlAzM4sZ5b1mNLD5gtgfJFOWYbTTet3t/sTvnZa15n5W9Tvqr1qXxRO6xz5Sfv+J21L9C+1iv0t/fbW9F+loLHK1T71tloVe1na8hydr1Pa+iqMm+54E4JX5bLZH4TE2UGyv6Spboq902/ZH27akDRAUzL3/vag74Tlb96kUGyCeFAGfHOAIzIzKNWuk5IAVjywYjAcEZ1z2cuEYEz4DiYE17hJrmidgRmDiBgb/F7wOHTPuRSzTnGi6EbA1EXULHtE8SwXMawInhvSBVl8nobGO76j6I6wrAIZlJt9p8LdKF8dgL9DNuPwVYVJGLdwVb0tt8Ztn8gbD8Un7e+kXvG/i9hX+8zZKdrnLFf3boCj9P3AA2PAs+424I7Q9D0bgt39Db/1wTJuXX+/x/Uyf8=</diagram></mxfile>
\ No newline at end of file diff --git a/tensorflow/core/BUILD b/tensorflow/core/BUILD index fc49825748..94973a0e52 100644 --- a/tensorflow/core/BUILD +++ b/tensorflow/core/BUILD @@ -279,6 +279,7 @@ cc_library( "platform/platform.h", "platform/protobuf.h", "platform/types.h", + "platform/windows/cpu_info.h", "lib/bfloat16/bfloat16.h", ] + tf_additional_proto_hdrs() + glob(tf_env_time_hdrs()), copts = tf_copts(), @@ -865,6 +866,7 @@ cc_library( "//tensorflow/core/kernels:mkl_pooling_ops", "//tensorflow/core/kernels:mkl_relu_op", "//tensorflow/core/kernels:mkl_reshape_op", + "//tensorflow/core/kernels:mkl_softmax_op", "//tensorflow/core/kernels:mkl_tfconv_op", "//tensorflow/core/kernels:mkl_aggregate_ops", ]), @@ -2831,6 +2833,7 @@ tf_cc_test_mkl( "//tensorflow/core/kernels:mkl_pooling_ops", "//tensorflow/core/kernels:mkl_relu_op", "//tensorflow/core/kernels:mkl_reshape_op", + "//tensorflow/core/kernels:mkl_softmax_op", "//tensorflow/core/kernels:mkl_tfconv_op", ]), ) diff --git a/tensorflow/core/framework/types.h b/tensorflow/core/framework/types.h index cb8e77f1df..ded6aa0991 100644 --- a/tensorflow/core/framework/types.h +++ b/tensorflow/core/framework/types.h @@ -453,6 +453,13 @@ inline bool DataTypeIsInteger(DataType dt) { return kDataTypeIsInteger.Contains(dt); } +// Is the dtype a signed integral type? +constexpr DataTypeSet kDataTypeIsSigned = + ToSet(DT_INT8) | ToSet(DT_INT16) | ToSet(DT_INT32) | ToSet(DT_INT64); +inline bool DataTypeIsSigned(DataType dt) { + return kDataTypeIsSigned.Contains(dt); +} + // Is the dtype an unsigned integral type? constexpr DataTypeSet kDataTypeIsUnsigned = ToSet(DT_UINT8) | ToSet(DT_UINT16) | ToSet(DT_UINT32) | ToSet(DT_UINT64); diff --git a/tensorflow/core/graph/mkl_layout_pass.cc b/tensorflow/core/graph/mkl_layout_pass.cc index 89b23f22fd..55bc401b9d 100644 --- a/tensorflow/core/graph/mkl_layout_pass.cc +++ b/tensorflow/core/graph/mkl_layout_pass.cc @@ -2456,9 +2456,9 @@ class MklLayoutRewritePass : public GraphOptimizationPass { // NOTE: names are alphabetically sorted. rinfo_.push_back({csinfo_.addn, mkl_op_registry::GetMklOpName(csinfo_.addn), CopyAttrsAddN, AddNRewrite}); - rinfo_.push_back({csinfo_.add, + /* rinfo_.push_back({csinfo_.add, mkl_op_registry::GetMklOpName(csinfo_.add), - CopyAttrsDataType, AlwaysRewrite}); + CopyAttrsDataType, AlwaysRewrite}); */ rinfo_.push_back({csinfo_.avg_pool, mkl_op_registry::GetMklOpName(csinfo_.avg_pool), CopyAttrsPooling, AlwaysRewrite}); @@ -3117,7 +3117,9 @@ void MklLayoutRewritePass::GetDummyMklTensorNode(std::unique_ptr<Graph>* g, Node* orig_input0 = nullptr; TF_CHECK_OK(orig_node->input_node(0, const_cast<const Node**>(&orig_input0))); - CHECK_NOTNULL((*g)->AddControlEdge(orig_input0, *out)); + // Allow duplicate while adding control edge as it would fail (return + // NULL) if we try to add duplicate edge. + CHECK_NOTNULL((*g)->AddControlEdge(orig_input0, *out, true)); } (*out)->set_assigned_device_name(orig_node->assigned_device_name()); @@ -3382,8 +3384,8 @@ void MklLayoutRewritePass::GetDummyWorkspaceTensorNode( std::unique_ptr<Graph>* g, Node** out, Node* orig_node) { // We use a tensor of shape {1} and value 0 to represent // dummy float tensor. We need this as a dummy workspace tensor. - // Workspace tensor has type float. - const DataType dt = DataTypeToEnum<float>::v(); + // Workspace tensor has type uint8. + const DataType dt = DataTypeToEnum<uint8>::v(); TensorProto proto; proto.set_dtype(dt); float zero[1] = {0}; @@ -3413,7 +3415,9 @@ void MklLayoutRewritePass::GetDummyWorkspaceTensorNode( Node* orig_input0 = nullptr; TF_CHECK_OK(orig_node->input_node(0, const_cast<const Node**>(&orig_input0))); - CHECK_NOTNULL((*g)->AddControlEdge(orig_input0, *out)); + // Allow duplicate while adding control edge as it would fail (return + // NULL) if we try to add duplicate edge. + CHECK_NOTNULL((*g)->AddControlEdge(orig_input0, *out, true)); } (*out)->set_assigned_device_name(orig_node->assigned_device_name()); @@ -3863,12 +3867,16 @@ Status MklLayoutRewritePass::MergeConv2DWithBiasAdd(std::unique_ptr<Graph>* g, // node are already copied in BuildNode. We handle control edges now. for (const Edge* e : pred->in_edges()) { if (e->IsControlEdge()) { - CHECK_NOTNULL((*g)->AddControlEdge(e->src(), new_node)); + // Allow duplicate while adding control edge as it would fail (return + // NULL) if we try to add duplicate edge. + CHECK_NOTNULL((*g)->AddControlEdge(e->src(), new_node, true)); } } for (const Edge* e : succ->in_edges()) { if (e->IsControlEdge()) { - CHECK_NOTNULL((*g)->AddControlEdge(e->src(), new_node)); + // Allow duplicate while adding control edge as it would fail (return + // NULL) if we try to add duplicate edge. + CHECK_NOTNULL((*g)->AddControlEdge(e->src(), new_node, true)); } } @@ -3876,14 +3884,18 @@ Status MklLayoutRewritePass::MergeConv2DWithBiasAdd(std::unique_ptr<Graph>* g, // First, we will fix outgoing control edges from 'pred' node. for (const Edge* e : pred->out_edges()) { if (e->IsControlEdge()) { - CHECK_NOTNULL((*g)->AddControlEdge(new_node, e->dst())); + // Allow duplicate while adding control edge as it would fail (return + // NULL) if we try to add duplicate edge. + CHECK_NOTNULL((*g)->AddControlEdge(new_node, e->dst(), true)); } } // Second, we will fix outgoing control and data edges from 'succ' node. for (const Edge* e : succ->out_edges()) { if (e->IsControlEdge()) { - CHECK_NOTNULL((*g)->AddControlEdge(new_node, e->dst())); + // Allow duplicate while adding control edge as it would fail (return + // NULL) if we try to add duplicate edge. + CHECK_NOTNULL((*g)->AddControlEdge(new_node, e->dst(), true)); } else { // BiasAdd has only 1 output (at slot 0) and merged node also has only 1 // output (at slot 0). @@ -3966,12 +3978,16 @@ Status MklLayoutRewritePass::MergeConv2DBackpropFilterWithBiasAddGrad( // edges now. for (const Edge* e : badd->in_edges()) { if (e->IsControlEdge()) { - CHECK_NOTNULL((*g)->AddControlEdge(e->src(), new_node)); + // Allow duplicate while adding control edge as it would fail (return + // NULL) if we try to add duplicate edge. + CHECK_NOTNULL((*g)->AddControlEdge(e->src(), new_node, true)); } } for (const Edge* e : fltr->in_edges()) { if (e->IsControlEdge()) { - CHECK_NOTNULL((*g)->AddControlEdge(e->src(), new_node)); + // Allow duplicate while adding control edge as it would fail (return + // NULL) if we try to add duplicate edge. + CHECK_NOTNULL((*g)->AddControlEdge(e->src(), new_node, true)); } } @@ -3987,7 +4003,9 @@ Status MklLayoutRewritePass::MergeConv2DBackpropFilterWithBiasAddGrad( for (const Edge* e : badd->out_edges()) { if (e->IsControlEdge()) { - CHECK_NOTNULL((*g)->AddControlEdge(new_node, e->dst())); + // Allow duplicate while adding control edge as it would fail (return + // NULL) if we try to add duplicate edge. + CHECK_NOTNULL((*g)->AddControlEdge(new_node, e->dst(), true)); } else { CHECK_NOTNULL((*g)->AddEdge(new_node, kMergedNodeBiasGradOutputIdx, e->dst(), e->dst_input())); @@ -3997,7 +4015,11 @@ Status MklLayoutRewritePass::MergeConv2DBackpropFilterWithBiasAddGrad( // Second, we will fix outgoing control and data edges from 'fltr' node. for (const Edge* e : fltr->out_edges()) { if (e->IsControlEdge()) { - CHECK_NOTNULL((*g)->AddControlEdge(new_node, e->dst())); + // We allow duplicate edge for this case since we already add control + // edge from new_node in line 3990. Line below could be adding same + // edge to same destination again. In such case, if we do not allow + // duplicate edge, then this call will fail. + CHECK_NOTNULL((*g)->AddControlEdge(new_node, e->dst(), true)); } else { CHECK_NOTNULL((*g)->AddEdge(new_node, kMergedNodeFilterGradOutputIdx, e->dst(), e->dst_input())); @@ -4091,7 +4113,9 @@ Status MklLayoutRewritePass::RewriteNode(std::unique_ptr<Graph>* g, // already copied in BuildNode. We need to handle control edges now. for (const Edge* e : orig_node->in_edges()) { if (e->IsControlEdge()) { - CHECK_NOTNULL((*g)->AddControlEdge(e->src(), new_node)); + // Allow duplicate while adding control edge as it would fail (return + // NULL) if we try to add duplicate edge. + CHECK_NOTNULL((*g)->AddControlEdge(e->src(), new_node, true)); } } @@ -4104,7 +4128,9 @@ Status MklLayoutRewritePass::RewriteNode(std::unique_ptr<Graph>* g, // GetTensorDataIndex provides this mapping function. for (const Edge* e : orig_node->out_edges()) { if (e->IsControlEdge()) { - CHECK_NOTNULL((*g)->AddControlEdge(new_node, e->dst())); + // Allow duplicate while adding control edge as it would fail (return + // NULL) if we try to add duplicate edge. + CHECK_NOTNULL((*g)->AddControlEdge(new_node, e->dst(), true)); } else { CHECK_NOTNULL((*g)->AddEdge(new_node, GetTensorDataIndex(e->src_output(), e->src()->num_outputs()), diff --git a/tensorflow/core/kernels/BUILD b/tensorflow/core/kernels/BUILD index 96f0feaf3d..fd99409c9b 100644 --- a/tensorflow/core/kernels/BUILD +++ b/tensorflow/core/kernels/BUILD @@ -5847,6 +5847,23 @@ tf_mkl_kernel_library( ) tf_mkl_kernel_library( + name = "mkl_softmax_op", + prefix = "mkl_softmax", + deps = [ + ":bounds_check", + ":ops_util", + "//tensorflow/core:core_cpu", + "//tensorflow/core:framework", + "//tensorflow/core:lib", + "//tensorflow/core:lib_internal", + "//tensorflow/core:nn_ops_op_lib", + ] + if_mkl([ + "//third_party/mkl:intel_binary_blob", + "@mkl_dnn//:mkl_dnn", + ]), +) + +tf_mkl_kernel_library( name = "mkl_fused_batch_norm_op", srcs = ["mkl_fused_batch_norm_op.cc"], deps = NN_DEPS + if_mkl([ diff --git a/tensorflow/core/kernels/cuda_solvers.h b/tensorflow/core/kernels/cuda_solvers.h index 3c389a82ab..ecfa23750c 100644 --- a/tensorflow/core/kernels/cuda_solvers.h +++ b/tensorflow/core/kernels/cuda_solvers.h @@ -427,7 +427,7 @@ inline DeviceLapackInfo CudaSolver::GetDeviceLapackInfo( int64 size, const string& debug_info) { DeviceLapackInfo new_dev_info(context_, size, debug_info); scratch_tensor_refs_.emplace_back(new_dev_info.tensor()); - return std::move(new_dev_info); + return new_dev_info; } } // namespace tensorflow diff --git a/tensorflow/core/kernels/cwise_op_pow.cc b/tensorflow/core/kernels/cwise_op_pow.cc index 5fb0735ac1..cf86478b0f 100644 --- a/tensorflow/core/kernels/cwise_op_pow.cc +++ b/tensorflow/core/kernels/cwise_op_pow.cc @@ -16,8 +16,9 @@ limitations under the License. #include "tensorflow/core/kernels/cwise_ops_common.h" namespace tensorflow { -REGISTER7(BinaryOp, CPU, "Pow", functor::pow, float, Eigen::half, double, int32, - int64, complex64, complex128); +REGISTER5(BinaryOp, CPU, "Pow", functor::pow, float, Eigen::half, double, + complex64, complex128); +REGISTER2(BinaryOp, CPU, "Pow", functor::safe_pow, int32, int64); #if GOOGLE_CUDA REGISTER4(BinaryOp, GPU, "Pow", functor::pow, float, Eigen::half, double, @@ -25,5 +26,5 @@ REGISTER4(BinaryOp, GPU, "Pow", functor::pow, float, Eigen::half, double, #endif #ifdef TENSORFLOW_USE_SYCL REGISTER2(BinaryOp, SYCL, "Pow", functor::pow, float, double); -#endif // TENSORFLOW_USE_SYCL +#endif // TENSORFLOW_USE_SYCL } // namespace tensorflow diff --git a/tensorflow/core/kernels/cwise_ops.h b/tensorflow/core/kernels/cwise_ops.h index da70b1e314..06918075a4 100644 --- a/tensorflow/core/kernels/cwise_ops.h +++ b/tensorflow/core/kernels/cwise_ops.h @@ -21,6 +21,7 @@ limitations under the License. #include <type_traits> #include "third_party/eigen3/unsupported/Eigen/CXX11/Tensor" + #include "tensorflow/core/framework/numeric_types.h" #include "tensorflow/core/framework/tensor_types.h" #include "tensorflow/core/kernels/bounds_check.h" @@ -115,6 +116,35 @@ struct functor_traits<scalar_binary_pow_op_google<Scalar, Exponent>> { enum { Cost = 5 * NumTraits<Scalar>::MulCost, PacketAccess = false }; }; +template <typename Scalar, typename Exponent> +struct safe_scalar_binary_pow_op { + static_assert(std::is_integral<Scalar>::value, "Integer type expected"); + static_assert(std::is_integral<Exponent>::value && + std::is_signed<Exponent>::value, + "Signed integer type expected"); + + bool* const error; + + EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE safe_scalar_binary_pow_op(bool* error) + : error(error) {} + + EIGEN_DEVICE_FUNC inline Scalar operator()(const Scalar& a, + const Exponent& b) const { + const Exponent safe_b = tensorflow::internal::SubtleMustCopy(b); + if (TF_PREDICT_TRUE(safe_b >= 0)) { + return numext::pow(a, safe_b); + } else { + *error = true; + return 0; + } + } +}; + +template <typename Scalar, typename Exponent> +struct functor_traits<safe_scalar_binary_pow_op<Scalar, Exponent>> { + enum { Cost = 5 * NumTraits<Scalar>::MulCost, PacketAccess = false }; +}; + template <typename T, typename DivOrMod> struct safe_div_or_mod_op { static_assert(std::is_integral<T>::value, "Integer type expected"); @@ -742,6 +772,11 @@ template <typename T> struct pow : base<T, Eigen::internal::scalar_binary_pow_op_google<T, T>> {}; template <typename T> +struct safe_pow : base<T, Eigen::internal::safe_scalar_binary_pow_op<T, T>> { + static const bool has_errors = true; +}; + +template <typename T> struct maximum : base<T, Eigen::internal::scalar_max_op<T>> {}; template <typename T> diff --git a/tensorflow/core/kernels/cwise_ops_common.cc b/tensorflow/core/kernels/cwise_ops_common.cc index 693c6467ac..e561e59cf5 100644 --- a/tensorflow/core/kernels/cwise_ops_common.cc +++ b/tensorflow/core/kernels/cwise_ops_common.cc @@ -40,6 +40,11 @@ void BinaryOpShared::SetComputeError(OpKernelContext* ctx) { if ((op == "Div" || op == "Mod" || op == "FloorMod" || op == "FloorDiv") && DataTypeIsInteger(ctx->op_kernel().input_type(0))) { ctx->CtxFailure(errors::InvalidArgument("Integer division by zero")); + } else if ((op == "Pow") && + DataTypeIsInteger(ctx->op_kernel().input_type(0)) && + DataTypeIsSigned(ctx->op_kernel().input_type(1))) { + ctx->CtxFailure(errors::InvalidArgument( + "Integers to negative integer powers are not allowed")); } else { ctx->CtxFailure( errors::Internal("Unexpected error in binary operator " diff --git a/tensorflow/core/kernels/decode_image_op.cc b/tensorflow/core/kernels/decode_image_op.cc index ceb152c3f0..44dcbf834c 100644 --- a/tensorflow/core/kernels/decode_image_op.cc +++ b/tensorflow/core/kernels/decode_image_op.cc @@ -87,11 +87,10 @@ class DecodeImageOp : public OpKernel { channels_ = 3; } else { OP_REQUIRES_OK(context, context->GetAttr("channels", &channels_)); - OP_REQUIRES( - context, - channels_ == 0 || channels_ == 1 || channels_ == 3 || channels_ == 4, - errors::InvalidArgument("channels must be 0, 1, 3, or 4, got ", - channels_)); + OP_REQUIRES(context, channels_ == 0 || channels_ == 1 || channels_ == 3 || + channels_ == 4, + errors::InvalidArgument( + "channels must be 0, 1, 3, or 4, got ", channels_)); } flags_.components = channels_; @@ -115,9 +114,8 @@ class DecodeImageOp : public OpKernel { if (format_ == kJpgFormat) { OP_REQUIRES_OK(context, context->GetAttr("ratio", &flags_.ratio)); - OP_REQUIRES(context, - flags_.ratio == 1 || flags_.ratio == 2 || flags_.ratio == 4 || - flags_.ratio == 8, + OP_REQUIRES(context, flags_.ratio == 1 || flags_.ratio == 2 || + flags_.ratio == 4 || flags_.ratio == 8, errors::InvalidArgument("ratio must be 1, 2, 4, or 8, got ", flags_.ratio)); OP_REQUIRES_OK(context, context->GetAttr("fancy_upscaling", @@ -132,9 +130,8 @@ class DecodeImageOp : public OpKernel { string dct_method; OP_REQUIRES_OK(context, context->GetAttr("dct_method", &dct_method)); OP_REQUIRES( - context, - (dct_method.empty() || dct_method == "INTEGER_FAST" || - dct_method == "INTEGER_ACCURATE"), + context, (dct_method.empty() || dct_method == "INTEGER_FAST" || + dct_method == "INTEGER_ACCURATE"), errors::InvalidArgument("dct_method must be one of " "{'', 'INTEGER_FAST', 'INTEGER_ACCURATE'}")); if (dct_method == "INTEGER_FAST") { @@ -160,9 +157,9 @@ class DecodeImageOp : public OpKernel { errors::InvalidArgument("Expected image (JPEG, PNG, or GIF), got ", FileFormatString(magic, input))); OP_REQUIRES(context, input.size() <= std::numeric_limits<int>::max(), - errors::InvalidArgument( - FileFormatString(magic, input), - " contents are too large for int: ", input.size())); + errors::InvalidArgument(FileFormatString(magic, input), + " contents are too large for int: ", + input.size())); OP_REQUIRES(context, magic == kPngFormat || channel_bits_ == 8, errors::InvalidArgument(FileFormatString(magic, input), " does not support uint16 output")); @@ -215,10 +212,9 @@ class DecodeImageOp : public OpKernel { input.data(), input.size(), flags, nullptr /* nwarn */, [=, &output](int width, int height, int channels) -> uint8* { Status status(context->allocate_output( - 0, - format_ == kGifFormat - ? TensorShape({1, height, width, channels}) - : TensorShape({height, width, channels}), + 0, format_ == kGifFormat + ? TensorShape({1, height, width, channels}) + : TensorShape({height, width, channels}), &output)); if (!status.ok()) { VLOG(1) << status; @@ -294,6 +290,7 @@ class DecodeImageOp : public OpKernel { // Decode GIF, allocating tensor once the size is known. Tensor* output = nullptr; + string error_string; OP_REQUIRES( context, gif::Decode(input.data(), input.size(), @@ -320,8 +317,10 @@ class DecodeImageOp : public OpKernel { return nullptr; } return output->flat<uint8>().data(); - }), - errors::InvalidArgument("Invalid GIF data, size ", input.size())); + }, + &error_string), + errors::InvalidArgument("Invalid GIF data (size ", input.size(), "), ", + error_string)); } private: diff --git a/tensorflow/core/kernels/matmul_op.cc b/tensorflow/core/kernels/matmul_op.cc index b3acb91202..cb68690f28 100644 --- a/tensorflow/core/kernels/matmul_op.cc +++ b/tensorflow/core/kernels/matmul_op.cc @@ -539,6 +539,7 @@ struct MatMulFunctor<SYCLDevice, T> { REGISTER_KERNEL_BUILDER( \ Name("MatMul").Device(DEVICE_CPU).TypeConstraint<T>("T").Label("eigen"), \ MatMulOp<CPUDevice, T, false /* cublas, ignored for CPU */>); + #define REGISTER_CPU(T) \ REGISTER_KERNEL_BUILDER( \ Name("MatMul").Device(DEVICE_CPU).TypeConstraint<T>("T"), \ diff --git a/tensorflow/core/kernels/mkl_aggregate_ops.cc b/tensorflow/core/kernels/mkl_aggregate_ops.cc index 44b94be3a0..bb5eceab27 100644 --- a/tensorflow/core/kernels/mkl_aggregate_ops.cc +++ b/tensorflow/core/kernels/mkl_aggregate_ops.cc @@ -61,6 +61,18 @@ class MklAddNOp : public OpKernel { GetMklShape(ctx, src2_idx, &(mkl_context.input2_shape)); bool input2_in_mkl_format = mkl_context.input2_shape.IsMklTensor(); + // if the shapes of two tensors are not same raise op error + TensorShape src1_shape, src2_shape; + src1_shape = input0.shape(); + src2_shape = input1.shape(); + if (!src1_shape.IsSameSize(src2_shape) ){ + ctx->SetStatus( + errors::InvalidArgument( + "Inputs to operation ", this->name(), " of type ", this->type_string(), + " must have the same size and shape. Input 0: ", + src1_shape.DebugString(), " != input 1: ", + src2_shape.DebugString())); + } // handle the case of a scalar if (!input1_in_mkl_format && input0.dims() == 0) { const TensorShape& o_shape = input0.shape(); @@ -307,6 +319,18 @@ class MklAddNOp : public OpKernel { src1_mkl_shape.GetDimension(): src1_tensor.dims(); int src2_dims_size = input2_in_mkl_format? src2_mkl_shape.GetDimension(): src2_tensor.dims(); + // if the shapes of two tensors are not same raise op error + TensorShape src1_shape, src2_shape; + src1_shape = src1_tensor.shape(); + src2_shape = src2_tensor.shape(); + if (!src1_shape.IsSameSize(src2_shape) ){ + ctx->SetStatus( + errors::InvalidArgument( + "Inputs to operation ", this->name(), " of type ", this->type_string(), + " must have the same size and shape. Input 0: ", + src1_shape.DebugString(), " != input 1: ", + src2_shape.DebugString())); + } if (!input1_in_mkl_format && src1_dims_size == 0) { Tensor* dst_tensor = nullptr; diff --git a/tensorflow/core/kernels/mkl_concat_op.cc b/tensorflow/core/kernels/mkl_concat_op.cc index 82771792d7..d109bb6bcf 100644 --- a/tensorflow/core/kernels/mkl_concat_op.cc +++ b/tensorflow/core/kernels/mkl_concat_op.cc @@ -598,7 +598,6 @@ class MklConcatOp : public OpKernel { concat_dim_tensor.shape().DebugString())); int32 concat_dim = internal::SubtleMustCopy( concat_dim_tensor.scalar<int32>()()); - if (concat_dim < 0) concat_dim = N + concat_dim; // check that ranks of all tensors match // and that their shapes match except for concat_dim. @@ -609,6 +608,9 @@ class MklConcatOp : public OpKernel { input_shapes[0].GetTfShape() : input_tensors[0].shape(); size_t expected_dims = expected_shape.dims(); + + if (concat_dim < 0) concat_dim = expected_dims + concat_dim; + for (auto& s : input_shapes) { if (s == expected_shape) {++i; continue;} diff --git a/tensorflow/core/kernels/mkl_conv_grad_filter_ops.cc b/tensorflow/core/kernels/mkl_conv_grad_filter_ops.cc index 793fa24d99..54d4916d49 100644 --- a/tensorflow/core/kernels/mkl_conv_grad_filter_ops.cc +++ b/tensorflow/core/kernels/mkl_conv_grad_filter_ops.cc @@ -467,6 +467,13 @@ class MklConv2DCustomBackpropFilterOp : return filter_tf_shape; } + TensorShape GetOutputTfShape(const TensorShape& input_shape, + const TensorShape& filter_shape, + const TensorShape& outbprop_shape) { + // Shape of output of Conv2DBackpropFilter is same as shape of filter. + return filter_shape; + } + const memory::dims& GetOutputDims(const memory::dims& fwd_input_dims, const memory::dims& fwd_filter_dims) { // Shape of output of Conv2DBackpropFilter is same as shape of filter. diff --git a/tensorflow/core/kernels/mkl_conv_grad_input_ops.cc b/tensorflow/core/kernels/mkl_conv_grad_input_ops.cc index db9e97e7ca..ef6db58d31 100644 --- a/tensorflow/core/kernels/mkl_conv_grad_input_ops.cc +++ b/tensorflow/core/kernels/mkl_conv_grad_input_ops.cc @@ -396,6 +396,13 @@ class MklConv2DCustomBackpropInputOp : return GetTfShape(context, kInputIndex_Filter); } + TensorShape GetOutputTfShape(const TensorShape& input_shape, + const TensorShape& filter_shape, + const TensorShape& outbprop_shape) { + // Output Shape of Conv2DBackpropInput is same as shape of Conv2D 'input'. + return input_shape; + } + const memory::dims& GetOutputDims(const memory::dims& fwd_input_dims, const memory::dims& fwd_filter_dims) { // Output Shape of Conv2DBackpropInput is same as shape of Conv2D 'input'. diff --git a/tensorflow/core/kernels/mkl_conv_ops.cc b/tensorflow/core/kernels/mkl_conv_ops.cc index a4e139bb54..0e77b45993 100644 --- a/tensorflow/core/kernels/mkl_conv_ops.cc +++ b/tensorflow/core/kernels/mkl_conv_ops.cc @@ -551,6 +551,13 @@ class MklConv2DOp : public OpKernel { output_mkl_shape.SetMklTensor(false); AllocateOutputSetMklShape(context, kOutputIndex_Dst, &output_tensor, src_tf_shape, output_mkl_shape); + + // MklConv2D also outputs converted filter as 2nd output of Conv2D. + filter_mkl_shape.SetMklTensor(false); + Tensor* output_filter_tensor = nullptr; + AllocateOutputSetMklShape(context, kOutputIndex_Filter, + &output_filter_tensor, + filter_tf_shape, filter_mkl_shape); return; } diff --git a/tensorflow/core/kernels/mkl_conv_ops.h b/tensorflow/core/kernels/mkl_conv_ops.h index b6883dbaa2..c6456bd5c3 100644 --- a/tensorflow/core/kernels/mkl_conv_ops.h +++ b/tensorflow/core/kernels/mkl_conv_ops.h @@ -390,6 +390,29 @@ class MklConv2DBackpropCommonOp : public OpKernel { TensorShape filter_tf_shape = MakeFilterTfShape(context, filter_tensor); TensorShape outbprop_tf_shape = GetTfShape(context, kOutbpropIdx); + // Corner cases: output with 0 elements and 0 batch size. + Tensor* output_tensor = nullptr; + if (input_tf_shape.num_elements() == 0 || + filter_tf_shape.num_elements() == 0 || + outbprop_tf_shape.num_elements() == 0) { + MklDnnShape output_mkl_shape; + output_mkl_shape.SetMklTensor(false); + TensorShape output_tf_shape = GetOutputTfShape(input_tf_shape, + filter_tf_shape, + outbprop_tf_shape); + const int kOutputIdx = 0; + AllocateOutputSetMklShape(context, kOutputIdx, &output_tensor, + output_tf_shape, output_mkl_shape); + CHECK_NOTNULL(output_tensor); + + // if output tensor has more than 0 elements, we need to 0 them out. + for (size_t i = 0; i < output_tf_shape.num_elements(); ++i) { + output_tensor->flat<T>().data()[i] = 0; + } + + return; + } + // By default, all dims are in MKL order. Only dims in TF order // are those with prefix tf_order. memory::dims outbprop_dims, fwd_input_dims, fwd_filter_dims; @@ -471,7 +494,6 @@ class MklConv2DBackpropCommonOp : public OpKernel { output.SetOpMemDesc(bwd_output_dims, memory::format::any); // Operator-specific call to create and execute primitive. - Tensor* output_tensor = nullptr; CreatePrimitive(context, cpu_engine, fwd_pd, &input, &filter, &outbackprop, &output, &output_tensor, strides, padding_l, padding_r, @@ -507,6 +529,11 @@ class MklConv2DBackpropCommonOp : public OpKernel { virtual TensorShape MakeFilterTfShape(OpKernelContext* context, const Tensor& filter_tensor) = 0; + /// Get the TensorFlow shape of output tensor. + virtual TensorShape GetOutputTfShape(const TensorShape& input_shape, + const TensorShape& filter_shape, + const TensorShape& outbprop_shape) = 0; + /// Get shape of output in MKL-DNN order. Computes shape of output from /// input shape (fwd_input_dims) and filter shape (fwd_filter_dims). virtual diff --git a/tensorflow/core/kernels/mkl_fused_batch_norm_op.cc b/tensorflow/core/kernels/mkl_fused_batch_norm_op.cc index a761562a4b..8340a91d05 100644 --- a/tensorflow/core/kernels/mkl_fused_batch_norm_op.cc +++ b/tensorflow/core/kernels/mkl_fused_batch_norm_op.cc @@ -703,27 +703,31 @@ class MklFusedBatchNormOp : public OpKernel { void Compute(OpKernelContext* context) override { try { auto cpu_engine = engine(engine::cpu, 0); - const size_t src_index = 0; // index of src input tensor - const size_t scale_index = 1; // index of scale tensor - const size_t shift_index = 2; // index of shift tensor - const size_t mean_index = 3; // index of est_mean tensor - const size_t var_index = 4; // index of est_variance tensor - - const Tensor& src_tensor = MklGetInput(context, src_index); - const Tensor& scale_tensor = MklGetInput(context, scale_index); - const Tensor& shift_tensor = MklGetInput(context, shift_index); - const Tensor& est_mean_tensor = MklGetInput(context, mean_index); - const Tensor& est_variance_tensor = MklGetInput(context, var_index); - + const size_t kSrcIndex = 0; // index of src input tensor + const size_t kScaleIndex = 1; // index of scale tensor + const size_t kShiftIndex = 2; // index of shift tensor + const size_t kMeanIndex = 3; // index of est_mean tensor + const size_t kVarianceIndex = 4; // index of est_variance tensor + + const Tensor& src_tensor = MklGetInput(context, kSrcIndex); + const Tensor& scale_tensor = MklGetInput(context, kScaleIndex); + const Tensor& shift_tensor = MklGetInput(context, kShiftIndex); + const Tensor& est_mean_tensor = MklGetInput(context, kMeanIndex); + const Tensor& est_variance_tensor = MklGetInput(context, + kVarianceIndex); + + TensorShape tf_shape_src; MklDnnShape dnn_shape_src; - GetMklShape(context, src_index, &dnn_shape_src); + GetMklShape(context, kSrcIndex, &dnn_shape_src); if (dnn_shape_src.IsMklTensor()) { + tf_shape_src = dnn_shape_src.GetTfShape(); OP_REQUIRES(context, dnn_shape_src.GetDimension() == 4, errors::InvalidArgument( "input must be 4-dimensional", src_tensor.shape().DebugString())); } else { + tf_shape_src = src_tensor.shape(); OP_REQUIRES(context, src_tensor.dims() == 4, errors::InvalidArgument( "input must be 4-dimensional", @@ -756,39 +760,35 @@ class MklFusedBatchNormOp : public OpKernel { est_variance_tensor.shape().DebugString())); } + // special case: input with 0 element and 0 batch size + Tensor* dst_tensor = nullptr; + if (tf_shape_src.num_elements() == 0) { + HandleEmptyInput(context, + tf_shape_src, + scale_tensor.shape(), + &dst_tensor); + return; + } + if (dnn_shape_src.IsMklTensor()) depth_ = dnn_shape_src.DimSize(MklDnnDims::Dim_C); else ExtractParams(context); // Indices of output tensors - const size_t dst_index = 0; - const size_t batch_mean_index = 1; - const size_t batch_variance_index = 2; - const size_t saved_mean_index = 3; - const size_t saved_variance_index = 4; + const size_t kDstIndex = 0; - // allocate batch mean output tensor + // allocate 4 output TF tensors Tensor* batch_mean_tensor = nullptr; - MklDnnShape mkl_shape_batch_mean; - mkl_shape_batch_mean.SetMklTensor(false); - AllocateOutputSetMklShape(context, - batch_mean_index, - &batch_mean_tensor, - scale_tensor.shape(), - mkl_shape_batch_mean); - CHECK_NOTNULL(batch_mean_tensor); - - // Batch variance Tensor* batch_variance_tensor = nullptr; - MklDnnShape mkl_shape_batch_variance; - mkl_shape_batch_variance.SetMklTensor(false); - AllocateOutputSetMklShape(context, - batch_variance_index, - &batch_variance_tensor, - scale_tensor.shape(), - mkl_shape_batch_variance); - CHECK_NOTNULL(batch_variance_tensor); + Tensor* saved_mean_tensor = nullptr; + Tensor* saved_variance_tensor = nullptr; + AllocateTFOutputs(context, + scale_tensor.shape(), + &batch_mean_tensor, + &batch_variance_tensor, + &saved_mean_tensor, + &saved_variance_tensor); if (is_training_) SetMeanVariance(*batch_mean_tensor, *batch_variance_tensor); @@ -844,26 +844,6 @@ class MklFusedBatchNormOp : public OpKernel { weights_data[k + depth_] = shift_tf[k]; } - // Mean and variance (without Bessel's correction) saved for backward - // computation to serve as pre-computed mean and variance. - Tensor* saved_mean_tensor = nullptr; - MklDnnShape mkl_shape_saved_mean; - mkl_shape_saved_mean.SetMklTensor(false); - AllocateOutputSetMklShape(context, saved_mean_index, - &saved_mean_tensor, - scale_tensor.shape(), - mkl_shape_saved_mean); - CHECK_NOTNULL(saved_mean_tensor); - - Tensor* saved_variance_tensor = nullptr; - MklDnnShape mkl_shape_saved_variance; - mkl_shape_saved_variance.SetMklTensor(false); - AllocateOutputSetMklShape(context, saved_variance_index, - &saved_variance_tensor, - scale_tensor.shape(), - mkl_shape_saved_variance); - CHECK_NOTNULL(saved_variance_tensor); - // set mean primitive auto mean_desc = memory::desc({1, depth_}, MklDnnType<T>(), @@ -902,7 +882,6 @@ class MklFusedBatchNormOp : public OpKernel { // allocate dst tensor MklDnnShape dnn_shape_dst; TensorShape tf_shape_dst; - Tensor* dst_tensor = nullptr; if (dnn_shape_src.IsMklTensor()) { dnn_shape_dst.SetMklTensor(true); auto dst_pd = bnrm_fwd_pd.dst_primitive_desc(); @@ -915,7 +894,7 @@ class MklFusedBatchNormOp : public OpKernel { dnn_shape_dst.SetMklTensor(false); tf_shape_dst = src_tensor.shape(); } - AllocateOutputSetMklShape(context, dst_index, &dst_tensor, + AllocateOutputSetMklShape(context, kDstIndex, &dst_tensor, tf_shape_dst, dnn_shape_dst); // Output of batchnorm has same shape as input. @@ -958,10 +937,8 @@ class MklFusedBatchNormOp : public OpKernel { size_t adjust_size = orig_size - 1; adjust_factor = (static_cast<float>(orig_size)) / adjust_size; } - T* batch_variance_data_tf = reinterpret_cast<T*>( - batch_variance_tensor->flat<T>().data()); for (int k=0; k < depth_; k++) - batch_variance_data_tf[k] = + batch_variance_tensor->flat<T>().data()[k] = (reinterpret_cast<T*>(variance_m.get_data_handle()))[k] * adjust_factor; } catch (mkldnn::error &e) { @@ -994,8 +971,100 @@ class MklFusedBatchNormOp : public OpKernel { variance_values_ = reinterpret_cast<T*>( const_cast<T*>(variance.flat<T>().data())); } -}; + void HandleEmptyInput(OpKernelContext* context, + TensorShape tf_shape_src, + TensorShape tf_shape_scale, + Tensor** dst_tensor) { + CHECK_NOTNULL(dst_tensor); + + const size_t kDstIndex = 0; + MklDnnShape dnn_shape_dst; + dnn_shape_dst.SetMklTensor(false); + AllocateOutputSetMklShape(context, kDstIndex, dst_tensor, + tf_shape_src, dnn_shape_dst); + CHECK_NOTNULL(*dst_tensor); + memset(const_cast<char*>((*dst_tensor)->tensor_data().data()), 0, + (*dst_tensor)->tensor_data().size()); + + Tensor* batch_mean_tensor = nullptr; + Tensor* batch_variance_tensor = nullptr; + Tensor* saved_mean_tensor = nullptr; + Tensor* saved_variance_tensor = nullptr; + AllocateTFOutputs(context, tf_shape_scale, + &batch_mean_tensor, + &batch_variance_tensor, + &saved_mean_tensor, + &saved_variance_tensor); + } + + void AllocateTFOutputs(OpKernelContext* context, + TensorShape tf_shape_scale, + Tensor** batch_mean_tensor, + Tensor** batch_variance_tensor, + Tensor** saved_mean_tensor, + Tensor** saved_variance_tensor) { + CHECK_NOTNULL(batch_mean_tensor); + CHECK_NOTNULL(batch_variance_tensor); + CHECK_NOTNULL(saved_mean_tensor); + CHECK_NOTNULL(saved_variance_tensor); + + const size_t kBatchMeanIndex = 1; + const size_t kBatchVarianceIndex = 2; + const size_t kSavedMeanIndex = 3; + const size_t kSavedVarianceIndex = 4; + + // allocate batch mean output tensor + MklDnnShape mkl_shape_batch_mean; + mkl_shape_batch_mean.SetMklTensor(false); + AllocateOutputSetMklShape(context, + kBatchMeanIndex, + batch_mean_tensor, + tf_shape_scale, + mkl_shape_batch_mean); + CHECK_NOTNULL(*batch_mean_tensor); + // set NAN mean value in case of empty input tensor + for (int k=0; k < tf_shape_scale.num_elements(); k++) + (*batch_mean_tensor)->flat<T>().data()[k] = NAN; + + // allocate batch variance output tensor + MklDnnShape mkl_shape_batch_variance; + mkl_shape_batch_variance.SetMklTensor(false); + AllocateOutputSetMklShape(context, + kBatchVarianceIndex, + batch_variance_tensor, + tf_shape_scale, + mkl_shape_batch_variance); + CHECK_NOTNULL(*batch_variance_tensor); + // set NAN variance value in case of empty input tensor + for (int k=0; k < tf_shape_scale.num_elements(); k++) + (*batch_variance_tensor)->flat<T>().data()[k] = NAN; + + // Mean and variance (without Bessel's correction) saved for backward + // computation to serve as pre-computed mean and variance. + MklDnnShape mkl_shape_saved_mean; + mkl_shape_saved_mean.SetMklTensor(false); + AllocateOutputSetMklShape(context, kSavedMeanIndex, + saved_mean_tensor, + tf_shape_scale, + mkl_shape_saved_mean); + CHECK_NOTNULL(*saved_mean_tensor); + // set NAN mean value in case of empty input tensor + for (int k=0; k < tf_shape_scale.num_elements(); k++) + (*saved_mean_tensor)->flat<T>().data()[k] = NAN; + + MklDnnShape mkl_shape_saved_variance; + mkl_shape_saved_variance.SetMklTensor(false); + AllocateOutputSetMklShape(context, kSavedVarianceIndex, + saved_variance_tensor, + tf_shape_scale, + mkl_shape_saved_variance); + CHECK_NOTNULL(*saved_variance_tensor); + // set NAN variance value in case of empty input tensor + for (int k=0; k < tf_shape_scale.num_elements(); k++) + (*saved_variance_tensor)->flat<T>().data()[k] = NAN; + } +}; template <typename Device, typename T> class MklFusedBatchNormGradOp : public OpKernel { @@ -1009,34 +1078,37 @@ class MklFusedBatchNormGradOp : public OpKernel { OP_REQUIRES_OK(context, context->GetAttr("data_format", &tensor_format)); OP_REQUIRES(context, FormatFromString(tensor_format, &tensor_format_), errors::InvalidArgument("Invalid data format")); + OP_REQUIRES_OK(context, context->GetAttr("is_training", &is_training_)); } void Compute(OpKernelContext* context) override { try { auto cpu_engine = engine(engine::cpu, 0); - - const size_t diff_dst_index = 0; // index of diff_dst tensor - const size_t src_index = 1; // index of src input tensor - const size_t scale_index = 2; // index of scale tensor - const size_t mean_index = 3; // index of saved_mean tensor - const size_t variance_index = 4; // index of saved_variance tensor - const Tensor& diff_dst_tensor = MklGetInput(context, diff_dst_index); - const Tensor& src_tensor = MklGetInput(context, src_index); - const Tensor& scale_tensor = MklGetInput(context, scale_index); - const Tensor& saved_mean_tensor = MklGetInput(context, mean_index); + const size_t kDiffDstIndex = 0; // index of diff_dst tensor + const size_t kSrcIndex = 1; // index of src input tensor + const size_t kScaleIndex = 2; // index of scale tensor + const size_t kMeanIndex = 3; // index of saved_mean tensor + const size_t kVarianceIndex = 4; // index of saved_variance tensor + const Tensor& diff_dst_tensor = MklGetInput(context, kDiffDstIndex); + const Tensor& src_tensor = MklGetInput(context, kSrcIndex); + const Tensor& scale_tensor = MklGetInput(context, kScaleIndex); + const Tensor& saved_mean_tensor = MklGetInput(context, kMeanIndex); const Tensor& saved_variance_tensor = MklGetInput(context, - variance_index); + kVarianceIndex); MklDnnShape dnn_shape_src, dnn_shape_diff_dst; - GetMklShape(context, src_index, &dnn_shape_src); - GetMklShape(context, diff_dst_index, &dnn_shape_diff_dst); + GetMklShape(context, kSrcIndex, &dnn_shape_src); + GetMklShape(context, kDiffDstIndex, &dnn_shape_diff_dst); + TensorShape tf_shape_src, tf_shape_diff_dst; if (dnn_shape_diff_dst.IsMklTensor()) { + tf_shape_diff_dst = dnn_shape_diff_dst.GetTfShape(); OP_REQUIRES(context, dnn_shape_diff_dst.GetDimension() == 4, errors::InvalidArgument( "input must be 4-dimensional", diff_dst_tensor.shape().DebugString())); } else { + tf_shape_diff_dst = diff_dst_tensor.shape(); OP_REQUIRES(context, diff_dst_tensor.dims() == 4, errors::InvalidArgument( "input must be 4-dimensional", @@ -1044,11 +1116,13 @@ class MklFusedBatchNormGradOp : public OpKernel { } if (dnn_shape_src.IsMklTensor()) { + tf_shape_src = dnn_shape_src.GetTfShape(); OP_REQUIRES(context, dnn_shape_src.GetDimension() == 4, errors::InvalidArgument( "input must be 4-dimensional", src_tensor.shape().DebugString())); } else { + tf_shape_src = src_tensor.shape(); OP_REQUIRES(context, src_tensor.dims() == 4, errors::InvalidArgument( "input must be 4-dimensional", @@ -1069,6 +1143,15 @@ class MklFusedBatchNormGradOp : public OpKernel { "saved variance must be 1-dimensional", saved_variance_tensor.shape().DebugString())); + Tensor* diff_src_tensor = nullptr; + if (tf_shape_src.num_elements() == 0 || + tf_shape_diff_dst.num_elements() == 0) { + HandleEmptyInput(context, tf_shape_src, + scale_tensor.shape(), + &diff_src_tensor); + return; + } + if (dnn_shape_src.IsMklTensor()) depth_ = dnn_shape_src.DimSize(MklDnnDims::Dim_C); else @@ -1165,25 +1248,21 @@ class MklFusedBatchNormGradOp : public OpKernel { auto diff_weights_m = memory(diff_weights_pd); auto bnrm_fwd_desc = batch_normalization_forward::desc( - prop_kind::forward_training, - src.GetUsrMemDesc(), - epsilon_, - use_scale_shift); + prop_kind::forward_training, + src.GetUsrMemDesc(), + epsilon_, + is_training_ ? use_scale_shift : + (use_scale_shift | use_global_stats)); auto bnrm_fwd_pd = batch_normalization_forward::primitive_desc( bnrm_fwd_desc, cpu_engine); // Indices of output tensors - const size_t diff_src_index = 0; // index of diff_src tensor - const size_t diff_scale_index = 1; // index of diff_scale tensor - const size_t diff_shift_index = 2; // index of diff_shift tensor - const size_t p1_index = 3; // index of 1st placeholder tensor - const size_t p2_index = 4; // index of 2nd placeholder tensor + const size_t kDiffSrcIndex = 0; // index of diff_src tensor // allocate diff_src tensor MklDnnShape dnn_shape_diff_src; TensorShape tf_shape_diff_src; - Tensor* diff_src_tensor = nullptr; if (dnn_shape_src.IsMklTensor()) { dnn_shape_diff_src.SetMklTensor(true); auto diff_src_pd = bnrm_fwd_pd.dst_primitive_desc(); @@ -1201,7 +1280,7 @@ class MklFusedBatchNormGradOp : public OpKernel { dnn_shape_diff_src.SetMklTensor(false); tf_shape_diff_src = src_tensor.shape(); } - AllocateOutputSetMklShape(context, diff_src_index, &diff_src_tensor, + AllocateOutputSetMklShape(context, kDiffSrcIndex, &diff_src_tensor, tf_shape_diff_src, dnn_shape_diff_src); diff_src.SetUsrMem(src_md, diff_src_tensor); @@ -1212,7 +1291,15 @@ class MklFusedBatchNormGradOp : public OpKernel { diff_src.GetUsrMemDesc(), src.GetUsrMemDesc(), epsilon_, - use_scale_shift); + /* for inference, specify use_global_stats + 1. on fwd prop, use mean and variance + provided as inputs + 2. on bwd prop, mean and variance are + considered as constants. Thus, + reduce the amout of MKL computations + */ + is_training_ ? use_scale_shift : + (use_scale_shift | use_global_stats)); auto bnrm_bwd_pd = batch_normalization_backward::primitive_desc( bnrm_bwd_desc, cpu_engine, @@ -1232,41 +1319,22 @@ class MklFusedBatchNormGradOp : public OpKernel { net.push_back(bnrm_bwd_op); stream(stream::kind::eager).submit(net).wait(); - // separate out scale and shift grad and copy to individual tensors - const TensorShape& tf_shape_scale_shift = scale_tensor.shape(); + // allocate 4 output TF tensors Tensor* diff_scale_tensor = nullptr; - MklDnnShape mkl_shape_diff_scale; - mkl_shape_diff_scale.SetMklTensor(false); - AllocateOutputSetMklShape(context, diff_scale_index, &diff_scale_tensor, - tf_shape_scale_shift, mkl_shape_diff_scale); - Tensor* diff_shift_tensor = nullptr; - MklDnnShape mkl_shape_diff_shift; - mkl_shape_diff_shift.SetMklTensor(false); - AllocateOutputSetMklShape(context, diff_shift_index, &diff_shift_tensor, - tf_shape_scale_shift, mkl_shape_diff_shift); + AllocateTFOutputs(context, scale_tensor.shape(), + &diff_scale_tensor, + &diff_shift_tensor); // copy data: diff_scale and diff_shift T* diff_weights_data_dnn = reinterpret_cast<T*> (diff_weights_m.get_data_handle()); - float* diff_scale_data_tf = const_cast<float*>( - static_cast<const float*>(diff_scale_tensor->flat<T>().data())); - float* diff_shift_data_tf = const_cast<float*>( - static_cast<const float*>(diff_shift_tensor->flat<T>().data())); for (int i = 0; i < depth_; i++) { - diff_scale_data_tf[i] = diff_weights_data_dnn[i]; - diff_shift_data_tf[i] = diff_weights_data_dnn[i + depth_]; + diff_scale_tensor->flat<T>().data()[i] = + diff_weights_data_dnn[i]; + diff_shift_tensor->flat<T>().data()[i] = + diff_weights_data_dnn[i + depth_]; } - - // Placeholders for estimated_mean and estimated_variance, which are - // used for inference and thus not needed here for gradient computation. - Tensor* p1_tensor = nullptr, *p2_tensor = nullptr; - MklDnnShape mkl_shape_p; - mkl_shape_p.SetMklTensor(false); - AllocateOutputSetMklShape(context, p1_index, &p1_tensor, - TensorShape({}), mkl_shape_p); - AllocateOutputSetMklShape(context, p2_index, &p2_tensor, - TensorShape({}), mkl_shape_p); } catch (mkldnn::error &e) { string error_msg = "Status: " + std::to_string(e.status) + ", message: " + string(e.message) + @@ -1282,12 +1350,74 @@ class MklFusedBatchNormGradOp : public OpKernel { T epsilon_; TensorFormat tensor_format_; int depth_; // batch normalization is done for per channel. + bool is_training_; void ExtractParams(OpKernelContext* context) { const Tensor& input = MklGetInput(context, 0); depth_ = static_cast<int>(GetTensorDim(input, tensor_format_, 'C')); } + void HandleEmptyInput(OpKernelContext* context, + TensorShape tf_shape_src, + TensorShape tf_shape_scale_shift, + Tensor** diff_src_tensor) { + const size_t kDiffSrcIndex = 0; + + MklDnnShape dnn_shape_diff_src; + dnn_shape_diff_src.SetMklTensor(false); + AllocateOutputSetMklShape(context, kDiffSrcIndex, diff_src_tensor, + tf_shape_src, dnn_shape_diff_src); + for (size_t i=0; i < (*diff_src_tensor)->shape().num_elements(); i++) + (*diff_src_tensor)->flat<T>().data()[i] = 0; + + Tensor* diff_scale_tensor = nullptr; + Tensor* diff_shift_tensor = nullptr; + AllocateTFOutputs(context, + tf_shape_scale_shift, + &diff_scale_tensor, + &diff_shift_tensor); + } + + void AllocateTFOutputs(OpKernelContext* context, + TensorShape tf_shape_scale_shift, + Tensor** diff_scale_tensor, + Tensor** diff_shift_tensor) { + CHECK_NOTNULL(diff_scale_tensor); + CHECK_NOTNULL(diff_shift_tensor); + + const size_t kDiffScaleIndex = 1; + const size_t kDiffShiftIndex = 2; + const size_t kP1Index = 3; + const size_t kP2Index = 4; + + // separate out scale and shift grad and copy to individual tensors + MklDnnShape mkl_shape_diff_scale; + mkl_shape_diff_scale.SetMklTensor(false); + AllocateOutputSetMklShape(context, kDiffScaleIndex, diff_scale_tensor, + tf_shape_scale_shift, mkl_shape_diff_scale); + CHECK_NOTNULL(*diff_scale_tensor); + for (size_t i=0; i < (*diff_scale_tensor)->shape().num_elements(); i++) + (*diff_scale_tensor)->flat<T>().data()[i] = 0; + + MklDnnShape mkl_shape_diff_shift; + mkl_shape_diff_shift.SetMklTensor(false); + AllocateOutputSetMklShape(context, kDiffShiftIndex, diff_shift_tensor, + tf_shape_scale_shift, mkl_shape_diff_shift); + CHECK_NOTNULL(*diff_shift_tensor); + for (size_t i=0; i < (*diff_shift_tensor)->shape().num_elements(); i++) + (*diff_shift_tensor)->flat<T>().data()[i] = 0; + + // Placeholders for estimated_mean and estimated_variance, which are + // used for inference and thus not needed here for gradient computation. + Tensor* p1_tensor = nullptr, *p2_tensor = nullptr; + MklDnnShape mkl_shape_p; + mkl_shape_p.SetMklTensor(false); + AllocateOutputSetMklShape(context, kP1Index, &p1_tensor, + TensorShape({}), mkl_shape_p); + AllocateOutputSetMklShape(context, kP2Index, &p2_tensor, + TensorShape({}), mkl_shape_p); + } + memory::dims GetMeanVarianceDims() { return memory::dims({1, depth_}); } diff --git a/tensorflow/core/kernels/mkl_input_conversion_op.cc b/tensorflow/core/kernels/mkl_input_conversion_op.cc index 001834b13b..4b5f7b8310 100644 --- a/tensorflow/core/kernels/mkl_input_conversion_op.cc +++ b/tensorflow/core/kernels/mkl_input_conversion_op.cc @@ -396,7 +396,7 @@ class MklInputConversionOp : public OpKernel { auto cpu_engine = engine(engine::cpu, 0); MklDnnData<T> tf_input(&cpu_engine); auto input_tf_md = mkl_output_mkl_shape.GetTfLayout(); - tf_input.SetUsrMem(input_tf_md, &tf_tensor); + tf_input.SetUsrMem(input_tf_md, tf_tensor); // Create reorder between tensorflow layout and Mkl layout. std::vector<primitive> net; diff --git a/tensorflow/core/kernels/mkl_lrn_op.cc b/tensorflow/core/kernels/mkl_lrn_op.cc index a8f28202f4..95e0404ba8 100644 --- a/tensorflow/core/kernels/mkl_lrn_op.cc +++ b/tensorflow/core/kernels/mkl_lrn_op.cc @@ -43,7 +43,7 @@ limitations under the License. using mkldnn::lrn_forward; using mkldnn::lrn_backward; using mkldnn::prop_kind; -using mkldnn::algorithm::lrn_across_channels; +using mkldnn::lrn_across_channels; using mkldnn::stream; #endif @@ -910,17 +910,23 @@ class MklLRNOp : public OpKernel { Eigen::Tensor<T, 2, Eigen::RowMajor> multiplier(depth, depth); GetBandMatrix<T>(depth, depth_radius_, &multiplier); - Tensor *output_dnn_data, *workspace; - MklDnnShape mkl_output_mkl_shape, mkl_workspace_mkl_shape; + Tensor *output_dnn_data = nullptr; + MklDnnShape mkl_output_mkl_shape; mkl_output_mkl_shape.SetMklTensor(false); mkl_output_mkl_shape.SetDimensions(4); AllocateOutputSetMklShape(context, kIdxOutput, &output_dnn_data, input.shape(), mkl_output_mkl_shape); + CHECK_NOTNULL(output_dnn_data); - mkl_workspace_mkl_shape.SetMklTensor(false); - mkl_workspace_mkl_shape.SetDimensions(4); - AllocateOutputSetMklShape(context, kIdxWorkspace, &workspace, - input.shape(), mkl_workspace_mkl_shape); + Tensor* workspace_tensor = nullptr; + MklDnnShape workspace_mkl_shape; + workspace_mkl_shape.SetMklTensor(false); + TensorShape workspace_tf_shape; + workspace_tf_shape.AddDim(0); + AllocateOutputSetMklShape(context, kIdxWorkspace, + &workspace_tensor, + workspace_tf_shape, workspace_mkl_shape); + CHECK_NOTNULL(workspace_tensor); auto out_shaped = output_dnn_data->shaped<T, 2>({nodes * batch, depth}); Eigen::array<DimPair, 1> dims = {{DimPair(1, 0)}}; @@ -1344,12 +1350,14 @@ class MklLRNGradOp : public OpKernel { errors::InvalidArgument("Output image must be 4-dimensional")); } - if (workspace_dnn_shape.IsMklTensor()) { - OP_REQUIRES(context, workspace_dnn_shape.IsMklTensor() == false, - errors::InvalidArgument("Workspace should not be MKL Tensor.")); - } else { - OP_REQUIRES(context, workspace_tensor.dims() == 1, - errors::InvalidArgument("Workspace must be 1-dimensional")); + if (workspace_enabled_) { + if (workspace_dnn_shape.IsMklTensor()) { + OP_REQUIRES(context, workspace_dnn_shape.IsMklTensor() == false, + errors::InvalidArgument("Workspace should not be MKL Tensor.")); + } else { + OP_REQUIRES(context, workspace_tensor.dims() == 1, + errors::InvalidArgument("Workspace must be 1-dimensional")); + } } } diff --git a/tensorflow/core/kernels/mkl_maxpooling_op.cc b/tensorflow/core/kernels/mkl_maxpooling_op.cc index de4d7d2e72..82c5229bab 100644 --- a/tensorflow/core/kernels/mkl_maxpooling_op.cc +++ b/tensorflow/core/kernels/mkl_maxpooling_op.cc @@ -517,7 +517,7 @@ class MklMaxPoolingOp : public MklPoolingForwardOpBase<T> { MklDnnData<T> dnn_data_input(&cpu_engine); MklDnnData<T> dnn_data_output(&cpu_engine); - MklDnnData<T> dnn_data_wksp(&cpu_engine); + MklDnnData<uint8> dnn_data_wksp(&cpu_engine); // initialize variables for the pooling op MklPoolParameters pool_params; @@ -588,16 +588,16 @@ class MklMaxPoolingOp : public MklPoolingForwardOpBase<T> { void AllocateWorkspaceTensor(OpKernelContext* context, const pooling_forward::primitive_desc& pool_fwd_prim_desc, - MklDnnData<T>* dnn_data_wksp) { + MklDnnData<uint8>* dnn_data_wksp) { CHECK_NOTNULL(dnn_data_wksp); Tensor* workspace_tensor = nullptr; memory::primitive_desc workspace_pd = pool_fwd_prim_desc.workspace_primitive_desc(); - size_t workspace_t_elems = this->GetNumTElements(workspace_pd); + size_t workspace_bytes = workspace_pd.get_size(); MklDnnShape workspace_mkl_shape; workspace_mkl_shape.SetMklTensor(false); TensorShape workspace_tf_shape; - workspace_tf_shape.AddDim(workspace_t_elems); + workspace_tf_shape.AddDim(workspace_bytes); AllocateOutputSetMklShape(context, kOutputTensorIndexWorkspace, &workspace_tensor, workspace_tf_shape, workspace_mkl_shape); @@ -651,7 +651,7 @@ class MklMaxPoolingGradOp : public MklPoolingBackwardOpBase<T> { if (!context->status().ok()) return; MklDnnData<T> grad_dnn_data(&cpu_engine); - MklDnnData<T> workspace_dnn_data(&cpu_engine); + MklDnnData<uint8> workspace_dnn_data(&cpu_engine); MklDnnData<T> output_dnn_data(&cpu_engine); Tensor* output_tensor = nullptr; MklPoolParameters pool_params; @@ -770,7 +770,7 @@ class MklMaxPoolingGradOp : public MklPoolingBackwardOpBase<T> { void ConfigureWorkspace(const Tensor& workspace_tensor, memory::primitive_desc workspace_pd, - MklDnnData<T> *workspace_dnn_data) { + MklDnnData<uint8> *workspace_dnn_data) { CHECK_NOTNULL(workspace_dnn_data); workspace_dnn_data->SetUsrMem(workspace_pd, &workspace_tensor); @@ -811,7 +811,7 @@ class MklMaxPoolingGradOp : public MklPoolingBackwardOpBase<T> { errors::InvalidArgument("Gradient must be " "4-dimensional")); } - if (this->workspace_enabled_){ + if (this->workspace_enabled_) { // The workspace should not be an MKL tensor OP_REQUIRES(context, workspace_mkl_shape.IsMklTensor() == false, errors::InvalidArgument("Workspace tensor should not" diff --git a/tensorflow/core/kernels/mkl_pooling_ops_common.h b/tensorflow/core/kernels/mkl_pooling_ops_common.h index d33e91a15d..b974b2c59a 100644 --- a/tensorflow/core/kernels/mkl_pooling_ops_common.h +++ b/tensorflow/core/kernels/mkl_pooling_ops_common.h @@ -231,7 +231,7 @@ class MklPoolingForwardOpBase : public MklPoolingOpBase<T> { const pooling_forward::primitive_desc& pool_fwd_desc, const MklDnnData<T>* src, MklDnnData<T>* dst, - MklDnnData<T>* wksp = nullptr) { + MklDnnData<uint8>* wksp = nullptr) { std::vector<primitive> net; // Create pooling primitive and add it to net @@ -307,7 +307,7 @@ class MklPoolingBackwardOpBase : public MklPoolingOpBase<T> { MklDnnData<T>* input_gradient_diff_dst, MklDnnData<T>* output_diff_src, const memory::primitive_desc& target_diff_dst_pd, - const MklDnnData<T>* workspace = nullptr) { + const MklDnnData<uint8>* workspace = nullptr) { std::vector<primitive> net; diff --git a/tensorflow/core/kernels/mkl_reshape_op.cc b/tensorflow/core/kernels/mkl_reshape_op.cc index 11c92ebdb4..b41e529357 100644 --- a/tensorflow/core/kernels/mkl_reshape_op.cc +++ b/tensorflow/core/kernels/mkl_reshape_op.cc @@ -256,11 +256,18 @@ class MklReshapeOp : public OpKernel { AllocateOutputSetMklShape(context, kOutputSlotIdx, &output_tensor, shape_to, mkl_shape_output); - // Insert reorder between Mkl layout and TensorFlow layout. + // Insert reorder between Mkl layout and TensorFlow layout if + // needed. If reorder is not needed but reshape is needed (since + // shape_from != shape_to), then we just copy input tensor to + // output tensor with target shape (we cannot forward Mkl layout + // in such case because shape has changed.) std::vector<primitive> net; - CHECK_EQ(dnn_data_input.CheckReorderToOpMem(output_tf_pd, - output_tensor, &net), true); - stream(stream::kind::eager).submit(net).wait(); + if (dnn_data_input.CheckReorderToOpMem(output_tf_pd, + output_tensor, &net)) { + stream(stream::kind::eager).submit(net).wait(); + } else { + output_tensor->CopyFrom(input_tensor, shape_to); + } return; } else { // If dimensions that are being expanded or collapsed are diff --git a/tensorflow/core/kernels/mkl_softmax_op.cc b/tensorflow/core/kernels/mkl_softmax_op.cc new file mode 100644 index 0000000000..896d562933 --- /dev/null +++ b/tensorflow/core/kernels/mkl_softmax_op.cc @@ -0,0 +1,163 @@ +/* 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. +==============================================================================*/ + +// See docs in ../ops/nn_ops.cc. +#ifdef INTEL_MKL +#ifdef INTEL_MKL_DNN + +#include "tensorflow/core/framework/numeric_op.h" +#include "tensorflow/core/framework/op_kernel.h" +#include "tensorflow/core/framework/register_types.h" +#include "tensorflow/core/framework/tensor.h" +#include "tensorflow/core/lib/core/errors.h" +#include "tensorflow/core/util/tensor_format.h" +#include "third_party/eigen3/unsupported/Eigen/CXX11/Tensor" + +#include "mkldnn.h" +#include "mkldnn_types.h" +#include "tensorflow/core/platform/default/logging.h" +#include "tensorflow/core/util/mkl_util.h" + +#include "mkldnn.hpp" +using mkldnn::stream; +using mkldnn::prop_kind; +using mkldnn::softmax_forward; + +namespace tensorflow { + +typedef Eigen::ThreadPoolDevice CPUDevice; + + + +template <typename Device, typename T> +class MklSoftmaxOp : public OpKernel { + public: + ~MklSoftmaxOp() {} + + explicit MklSoftmaxOp(OpKernelConstruction* context) : OpKernel(context) {} + + void Compute(OpKernelContext* context) override { + try { + auto cpu_engine = engine(engine::cpu, 0); + + // src_tensor now points to the 0-th input of global data struct "context" + size_t src_idx = 0; + const Tensor& src_tensor = MklGetInput(context, src_idx); + + // Add: get MklShape + MklDnnShape src_mkl_shape; + GetMklShape(context, src_idx, &src_mkl_shape); + + + // src_dims is the dimenstion of src_tensor + // dim of the dst will also be same as src_dims + auto src_tf_shape = src_mkl_shape.IsMklTensor() ? + src_mkl_shape.GetTfShape() : src_tensor.shape(); + auto src_dims = TFShapeToMklDnnDims(src_tf_shape); + auto output_dims = src_dims; + + // Create softmax memory for src, dst: both are defined in mkl_util.h, + // they are wrapper + MklDnnData<T> src(&cpu_engine); + MklDnnData<T> dst(&cpu_engine); + + // If input is in MKL layout, then simply grab input layout; otherwise, + // construct input Tf layout. For TF layout, although input shape + // (src_dims) required is in MKL-DNN order, the layout is Tensorflow's + // layout + auto src_md = src_mkl_shape.IsMklTensor() + ? src_mkl_shape.GetMklLayout() + : memory::desc(src_dims, MklDnnType<T>(), + memory::format::nc); + + // src: setting memory descriptor and op memory descriptor + // Basically following two functions maps the TF "src_tensor" to mkl + // tensor object "src" + // following functions are in mkl_util.h + // data format is "nc" for src and dst; since the src and dst buffer is + // always in 2D shape + src.SetUsrMem(src_md, &src_tensor); + src.SetOpMemDesc(src_dims, memory::format::nc); + + // creating a memory descriptor + int axis = 1; // axis to which softmax will be applied + auto softmax_fwd_desc = softmax_forward::desc(prop_kind::forward_scoring, + src.GetOpMemDesc(), axis); + auto softmax_fwd_pd = softmax_forward::primitive_desc(softmax_fwd_desc, + cpu_engine); + + // add: output + Tensor* output_tensor = nullptr; + MklDnnShape output_mkl_shape; + TensorShape output_tf_shape; // shape of output TF tensor. + // Softmax MklDnn output layout is same as input layout. + auto dst_pd = src.GetUsrMemPrimDesc(); + + // if input is MKL shape, ouput is also MKL shape. + // if input is TF shape, output is also TF shape + if (src_mkl_shape.IsMklTensor()) { + output_mkl_shape.SetMklTensor(true); + output_mkl_shape.SetMklLayout(&dst_pd); + output_mkl_shape.SetElemType(MklDnnType<T>()); + output_mkl_shape.SetTfLayout(output_dims.size(), output_dims, + memory::format::nc); + output_tf_shape.AddDim((dst_pd.get_size() / sizeof(T))); + } else { // then output is also TF shape + output_mkl_shape.SetMklTensor(false); + output_tf_shape = MklDnnDimsToTFShape(output_dims); + } + // Allocate output shape (MKL or TF based on the above) + AllocateOutputSetMklShape(context, 0, &output_tensor, output_tf_shape, + output_mkl_shape); + + // Output_dims and input_dims are same + dst.SetUsrMem(src_md, output_tensor); + + // finally creating the "softmax op" using the primitive descriptor, src + // and dst + auto softmax_fwd = + softmax_forward(softmax_fwd_pd, src.GetOpMem(), dst.GetOpMem()); + + // execute net (pushing to the stream) + // following 3 are common for all mkl dnn ops + std::vector<primitive> net; + net.push_back(softmax_fwd); + stream(stream::kind::eager).submit(net).wait(); + } catch (mkldnn::error& e) { + string error_msg = "Status: " + std::to_string(e.status) + ", message: " + + string(e.message) + ", in file " + string(__FILE__) + + ":" + std::to_string(__LINE__); + OP_REQUIRES_OK( + context, + errors::Aborted("Operation received an exception:", error_msg)); + } + } +}; + +/* Register DNN kernels for supported operations and supported types - right now + * it is only Softmax and f32 */ +#define REGISTER_SOFTMAX_MKL_SUPPORTED_KERNELS_TYPES(type) \ + REGISTER_KERNEL_BUILDER(Name("_MklSoftmax") \ + .Device(DEVICE_CPU) \ + .TypeConstraint<type>("T") \ + .Label(mkl_op_registry::kMklOpLabel), \ + MklSoftmaxOp<CPUDevice, type>); +TF_CALL_float(REGISTER_SOFTMAX_MKL_SUPPORTED_KERNELS_TYPES); + + +} // namespace tensorflow + +#endif // INTEL_MKL_DNN +#endif // INTEL_MKL diff --git a/tensorflow/core/kernels/pooling_ops_common.cc b/tensorflow/core/kernels/pooling_ops_common.cc index 6a52a15c93..d4241b5809 100644 --- a/tensorflow/core/kernels/pooling_ops_common.cc +++ b/tensorflow/core/kernels/pooling_ops_common.cc @@ -222,7 +222,7 @@ void DnnPoolingOp<T>::Compute( output_desc, &output_data) .ok(); OP_REQUIRES(context, status, - errors::Internal("cudnn PoolBackward launch failed")); + errors::Internal("cudnn PoolForward launch failed")); if (data_format == FORMAT_NHWC) { /// Transform the output data from NCHW back to NHWC diff --git a/tensorflow/core/kernels/spectrogram_test_utils.cc b/tensorflow/core/kernels/spectrogram_test_utils.cc index 046f6344df..bc30330d61 100644 --- a/tensorflow/core/kernels/spectrogram_test_utils.cc +++ b/tensorflow/core/kernels/spectrogram_test_utils.cc @@ -70,10 +70,24 @@ bool ReadRawFloatFileToComplexVector( int offset = 0; const int end = data_string.size(); while (offset < end) { +#if __BYTE_ORDER__ == __ORDER_BIG_ENDIAN__ + char arr[4]; + for (int i = 0; i < kBytesPerValue; ++i ) { + arr[3 - i] = *(data_string.data() + offset + i); + } + memcpy(&real_out, arr, kBytesPerValue); + offset += kBytesPerValue; + for (int i = 0; i < kBytesPerValue; ++i ) { + arr[3 - i] = *(data_string.data() + offset + i); + } + memcpy(&imag_out, arr, kBytesPerValue); + offset += kBytesPerValue; +#else memcpy(&real_out, data_string.data() + offset, kBytesPerValue); offset += kBytesPerValue; memcpy(&imag_out, data_string.data() + offset, kBytesPerValue); offset += kBytesPerValue; +#endif if (row_counter >= row_length) { data->push_back(data_row); data_row.clear(); diff --git a/tensorflow/core/kernels/transpose_functor_cpu.cc b/tensorflow/core/kernels/transpose_functor_cpu.cc index 41b73fdaf4..6594f7ee7b 100644 --- a/tensorflow/core/kernels/transpose_functor_cpu.cc +++ b/tensorflow/core/kernels/transpose_functor_cpu.cc @@ -88,6 +88,18 @@ struct Transpose<CPUDevice, T, conjugate> { internal::TransposeUsingEigen<CPUDevice, T, 5>(d, in, perm, conjugate, out); break; + case 6: + internal::TransposeUsingEigen<CPUDevice, T, 6>(d, in, perm, conjugate, + out); + break; + case 7: + internal::TransposeUsingEigen<CPUDevice, T, 7>(d, in, perm, conjugate, + out); + break; + case 8: + internal::TransposeUsingEigen<CPUDevice, T, 8>(d, in, perm, conjugate, + out); + break; default: TransposeSimple<T, conjugate>(d, in, perm, out); break; diff --git a/tensorflow/core/kernels/transpose_functor_gpu.cu.cc b/tensorflow/core/kernels/transpose_functor_gpu.cu.cc index 493dac9a7c..d6a237d6c1 100644 --- a/tensorflow/core/kernels/transpose_functor_gpu.cu.cc +++ b/tensorflow/core/kernels/transpose_functor_gpu.cu.cc @@ -201,6 +201,27 @@ struct Transpose<GPUDevice, T, conjugate> { out); } break; + case 6: + if (!internal::TransposeUsingTile<T, conjugate>::run(d, in, perm, + out)) { + internal::TransposeUsingEigen<GPUDevice, T, 6>(d, in, perm, conjugate, + out); + } + break; + case 7: + if (!internal::TransposeUsingTile<T, conjugate>::run(d, in, perm, + out)) { + internal::TransposeUsingEigen<GPUDevice, T, 7>(d, in, perm, conjugate, + out); + } + break; + case 8: + if (!internal::TransposeUsingTile<T, conjugate>::run(d, in, perm, + out)) { + internal::TransposeUsingEigen<GPUDevice, T, 8>(d, in, perm, conjugate, + out); + } + break; default: internal::TransposeSimple<T, conjugate>(d, in, perm, out); break; diff --git a/tensorflow/core/kernels/xent_op.cc b/tensorflow/core/kernels/xent_op.cc index dc21cee3a8..0f8d027caa 100644 --- a/tensorflow/core/kernels/xent_op.cc +++ b/tensorflow/core/kernels/xent_op.cc @@ -67,10 +67,12 @@ class SoftmaxXentWithLogitsOp : public OpKernel { // Try to reuse the logits_in buffer for the backprop output. OP_REQUIRES_OK(context, context->forward_input_or_allocate_output( {0}, 1, logits_in.shape(), &back_out)); - functor::XentFunctor<Device, T> functor; - functor(context->eigen_device<Device>(), logits_in.matrix<T>(), - labels_in.matrix<T>(), scratch.matrix<T>(), loss_out->vec<T>(), - back_out->matrix<T>()); + if (logits_in.dim_size(0) > 0) { + functor::XentFunctor<Device, T> functor; + functor(context->eigen_device<Device>(), logits_in.matrix<T>(), + labels_in.matrix<T>(), scratch.matrix<T>(), loss_out->vec<T>(), + back_out->matrix<T>()); + } } }; diff --git a/tensorflow/core/lib/gif/gif_io.cc b/tensorflow/core/lib/gif/gif_io.cc index b5c0d9f621..0f6999c88f 100644 --- a/tensorflow/core/lib/gif/gif_io.cc +++ b/tensorflow/core/lib/gif/gif_io.cc @@ -17,6 +17,7 @@ limitations under the License. #include "tensorflow/core/lib/gif/gif_io.h" #include "tensorflow/core/lib/gtl/cleanup.h" +#include "tensorflow/core/lib/strings/strcat.h" #include "tensorflow/core/platform/gif.h" #include "tensorflow/core/platform/logging.h" #include "tensorflow/core/platform/mem.h" @@ -44,7 +45,8 @@ int input_callback(GifFileType* gif_file, GifByteType* buf, int size) { } uint8* Decode(const void* srcdata, int datasize, - std::function<uint8*(int, int, int, int)> allocate_output) { + const std::function<uint8*(int, int, int, int)>& allocate_output, + string* error_string) { int error_code = D_GIF_SUCCEEDED; InputBufferInfo info = {reinterpret_cast<const uint8*>(srcdata), datasize}; GifFileType* gif_file = @@ -57,17 +59,17 @@ uint8* Decode(const void* srcdata, int datasize, } }); if (error_code != D_GIF_SUCCEEDED) { - LOG(ERROR) << "Fail to open gif file, reason: " - << GifErrorString(error_code); + *error_string = strings::StrCat("failed to open gif file: ", + GifErrorString(error_code)); return nullptr; } if (DGifSlurp(gif_file) != GIF_OK) { - LOG(ERROR) << "Fail to slurp gif file, reason: " - << GifErrorString(gif_file->Error); + *error_string = strings::StrCat("failed to slurp gif file: ", + GifErrorString(gif_file->Error)); return nullptr; } if (gif_file->ImageCount <= 0) { - LOG(ERROR) << "Gif file does not contain any image"; + *error_string = strings::StrCat("gif file does not contain any image"); return nullptr; } @@ -83,7 +85,7 @@ uint8* Decode(const void* srcdata, int datasize, GifImageDesc* img_desc = &this_image->ImageDesc; if (img_desc->Left != 0 || img_desc->Top != 0 || img_desc->Width != width || img_desc->Height != height) { - LOG(ERROR) << "Can't process optimized gif."; + *error_string = strings::StrCat("can't process optimized gif"); return nullptr; } diff --git a/tensorflow/core/lib/gif/gif_io.h b/tensorflow/core/lib/gif/gif_io.h index 5399e6a538..0a7967a5a1 100644 --- a/tensorflow/core/lib/gif/gif_io.h +++ b/tensorflow/core/lib/gif/gif_io.h @@ -43,7 +43,8 @@ namespace tensorflow { namespace gif { uint8* Decode(const void* srcdata, int datasize, - std::function<uint8*(int, int, int, int)> allocate_output); + const std::function<uint8*(int, int, int, int)>& allocate_output, + string* error_string); } // namespace gif } // namespace tensorflow diff --git a/tensorflow/core/ops/nn_ops.cc b/tensorflow/core/ops/nn_ops.cc index 536fc7c0c1..3f72b41569 100644 --- a/tensorflow/core/ops/nn_ops.cc +++ b/tensorflow/core/ops/nn_ops.cc @@ -1818,7 +1818,11 @@ REGISTER_OP("_MklMaxPool") .Input("input: T") .Input("mkl_input: uint8") .Output("output: T") +#ifndef INTEL_MKL_DNN .Output("workspace: T") +#else + .Output("workspace: uint8") +#endif .Output("mkl_output: uint8") .Output("mkl_workspace: uint8") .SetShapeFn(shape_inference::MaxPoolShape) @@ -1840,7 +1844,11 @@ REGISTER_OP("_MklMaxPoolGrad") .Input("orig_input: T") .Input("orig_output: T") .Input("grad: T") +#ifndef INTEL_MKL_DNN .Input("workspace: T") +#else + .Input("workspace: uint8") +#endif .Input("mkl_orig_input: uint8") .Input("mkl_orig_output: uint8") .Input("mkl_grad: uint8") diff --git a/tensorflow/core/platform/s3/aws_logging.cc b/tensorflow/core/platform/s3/aws_logging.cc index 41b854d634..fbca0acc36 100644 --- a/tensorflow/core/platform/s3/aws_logging.cc +++ b/tensorflow/core/platform/s3/aws_logging.cc @@ -48,6 +48,7 @@ void AWSLogSystem::LogStream(Aws::Utils::Logging::LogLevel log_level, void AWSLogSystem::LogMessage(Aws::Utils::Logging::LogLevel log_level, const std::string& message) { + if (message == "Initializing Curl library") return; switch (log_level) { case Aws::Utils::Logging::LogLevel::Info: LOG(INFO) << message; diff --git a/tensorflow/core/platform/s3/s3_file_system.cc b/tensorflow/core/platform/s3/s3_file_system.cc index 397f26ec0b..ebda3a2065 100644 --- a/tensorflow/core/platform/s3/s3_file_system.cc +++ b/tensorflow/core/platform/s3/s3_file_system.cc @@ -14,11 +14,13 @@ limitations under the License. ==============================================================================*/ #include "tensorflow/core/platform/s3/s3_file_system.h" #include "tensorflow/core/lib/io/path.h" +#include "tensorflow/core/lib/strings/str_util.h" #include "tensorflow/core/platform/mutex.h" #include "tensorflow/core/platform/s3/aws_logging.h" #include "tensorflow/core/platform/s3/s3_crypto.h" #include <aws/core/Aws.h> +#include <aws/core/config/AWSProfileConfigLoader.h> #include <aws/core/utils/FileSystemUtils.h> #include <aws/core/utils/logging/AWSLogging.h> #include <aws/core/utils/logging/LogSystemInterface.h> @@ -54,13 +56,37 @@ Aws::Client::ClientConfiguration& GetDefaultClientConfig() { cfg.endpointOverride = Aws::String(endpoint); } const char* region = getenv("AWS_REGION"); + if (!region) { + // TODO (yongtang): `S3_REGION` should be deprecated after 2.0. + region = getenv("S3_REGION"); + } if (region) { cfg.region = Aws::String(region); } else { - // TODO (yongtang): `S3_REGION` should be deprecated after 2.0. - const char* region = getenv("S3_REGION"); - if (region) { - cfg.region = Aws::String(region); + // Load config file (e.g., ~/.aws/config) only if AWS_SDK_LOAD_CONFIG + // is set with a truthy value. + const char* load_config_env = getenv("AWS_SDK_LOAD_CONFIG"); + string load_config = + load_config_env ? str_util::Lowercase(load_config_env) : ""; + if (load_config == "true" || load_config == "1") { + Aws::String config_file; + // If AWS_CONFIG_FILE is set then use it, otherwise use ~/.aws/config. + const char* config_file_env = getenv("AWS_CONFIG_FILE"); + if (config_file_env) { + config_file = config_file_env; + } else { + const char* home_env = getenv("HOME"); + if (home_env) { + config_file = home_env; + config_file += "/.aws/config"; + } + } + Aws::Config::AWSConfigFileProfileConfigLoader loader(config_file); + loader.Load(); + auto profiles = loader.GetProfiles(); + if (!profiles["default"].GetRegion().empty()) { + cfg.region = profiles["default"].GetRegion(); + } } } const char* use_https = getenv("S3_USE_HTTPS"); @@ -79,6 +105,22 @@ Aws::Client::ClientConfiguration& GetDefaultClientConfig() { cfg.verifySSL = true; } } + const char* connect_timeout = getenv("S3_CONNECT_TIMEOUT_MSEC"); + if (connect_timeout) { + int64 timeout; + + if (strings::safe_strto64(connect_timeout, &timeout)) { + cfg.connectTimeoutMs = timeout; + } + } + const char* request_timeout = getenv("S3_REQUEST_TIMEOUT_MSEC"); + if (request_timeout) { + int64 timeout; + + if (strings::safe_strto64(request_timeout, &timeout)) { + cfg.requestTimeoutMs = timeout; + } + } init = true; } diff --git a/tensorflow/core/public/version.h b/tensorflow/core/public/version.h index 836efc079f..67da7bf452 100644 --- a/tensorflow/core/public/version.h +++ b/tensorflow/core/public/version.h @@ -24,7 +24,7 @@ limitations under the License. // TF_VERSION_SUFFIX is non-empty for pre-releases (e.g. "-alpha", "-alpha.1", // "-beta", "-rc", "-rc.1") -#define TF_VERSION_SUFFIX "-rc0" +#define TF_VERSION_SUFFIX "-rc1" #define TF_STR_HELPER(x) #x #define TF_STR(x) TF_STR_HELPER(x) diff --git a/tensorflow/docs_src/api_guides/python/python_io.md b/tensorflow/docs_src/api_guides/python/python_io.md index a5444408fe..06282e49d5 100644 --- a/tensorflow/docs_src/api_guides/python/python_io.md +++ b/tensorflow/docs_src/api_guides/python/python_io.md @@ -14,16 +14,16 @@ suitable if fast sharding or other non-sequential access is desired. ## TFRecords Format Details -A TFRecords file contains a sequence of strings with CRC hashes. Each record -has the format +A TFRecords file contains a sequence of strings with CRC32C (32-bit CRC using +the Castagnoli polynomial) hashes. Each record has the format uint64 length uint32 masked_crc32_of_length byte data[length] uint32 masked_crc32_of_data -and the records are concatenated together to produce the file. The CRC32s -are [described here](https://en.wikipedia.org/wiki/Cyclic_redundancy_check), -and the mask of a CRC is +and the records are concatenated together to produce the file. CRCs are +[described here](https://en.wikipedia.org/wiki/Cyclic_redundancy_check), and +the mask of a CRC is masked_crc = ((crc >> 15) | (crc << 17)) + 0xa282ead8ul diff --git a/tensorflow/docs_src/install/install_c.md b/tensorflow/docs_src/install/install_c.md index d79cd1415d..ba1a4118ae 100644 --- a/tensorflow/docs_src/install/install_c.md +++ b/tensorflow/docs_src/install/install_c.md @@ -38,7 +38,7 @@ enable TensorFlow for C: OS="linux" # Change to "darwin" for macOS TARGET_DIRECTORY="/usr/local" curl -L \ - "https://storage.googleapis.com/tensorflow/libtensorflow/libtensorflow-${TF_TYPE}-${OS}-x86_64-1.5.0-rc0.tar.gz" | + "https://storage.googleapis.com/tensorflow/libtensorflow/libtensorflow-${TF_TYPE}-${OS}-x86_64-1.5.0-rc1.tar.gz" | sudo tar -C $TARGET_DIRECTORY -xz The `tar` command extracts the TensorFlow C library into the `lib` diff --git a/tensorflow/docs_src/install/install_go.md b/tensorflow/docs_src/install/install_go.md index 49f5350405..87cc647317 100644 --- a/tensorflow/docs_src/install/install_go.md +++ b/tensorflow/docs_src/install/install_go.md @@ -38,7 +38,7 @@ steps to install this library and enable TensorFlow for Go: TF_TYPE="cpu" # Change to "gpu" for GPU support TARGET_DIRECTORY='/usr/local' curl -L \ - "https://storage.googleapis.com/tensorflow/libtensorflow/libtensorflow-${TF_TYPE}-$(go env GOOS)-x86_64-1.5.0-rc0.tar.gz" | + "https://storage.googleapis.com/tensorflow/libtensorflow/libtensorflow-${TF_TYPE}-$(go env GOOS)-x86_64-1.5.0-rc1.tar.gz" | sudo tar -C $TARGET_DIRECTORY -xz The `tar` command extracts the TensorFlow C library into the `lib` diff --git a/tensorflow/docs_src/install/install_java.md b/tensorflow/docs_src/install/install_java.md index 47b1251427..37e109a6e4 100644 --- a/tensorflow/docs_src/install/install_java.md +++ b/tensorflow/docs_src/install/install_java.md @@ -36,7 +36,7 @@ following to the project's `pom.xml` to use the TensorFlow Java APIs: <dependency> <groupId>org.tensorflow</groupId> <artifactId>tensorflow</artifactId> - <version>1.5.0-rc0</version> + <version>1.5.0-rc1</version> </dependency> ``` @@ -65,7 +65,7 @@ As an example, these steps will create a Maven project that uses TensorFlow: <dependency> <groupId>org.tensorflow</groupId> <artifactId>tensorflow</artifactId> - <version>1.5.0-rc0</version> + <version>1.5.0-rc1</version> </dependency> </dependencies> </project> @@ -123,12 +123,12 @@ instead: <dependency> <groupId>org.tensorflow</groupId> <artifactId>libtensorflow</artifactId> - <version>1.4.0</version> + <version>1.5.0-rc1</version> </dependency> <dependency> <groupId>org.tensorflow</groupId> <artifactId>libtensorflow_jni_gpu</artifactId> - <version>1.4.0</version> + <version>1.5.0-rc1</version> </dependency> ``` @@ -147,7 +147,7 @@ refer to the simpler instructions above instead. Take the following steps to install TensorFlow for Java on Linux or macOS: 1. Download - [libtensorflow.jar](https://storage.googleapis.com/tensorflow/libtensorflow/libtensorflow-1.5.0-rc0.jar), + [libtensorflow.jar](https://storage.googleapis.com/tensorflow/libtensorflow/libtensorflow-1.5.0-rc1.jar), which is the TensorFlow Java Archive (JAR). 2. Decide whether you will run TensorFlow for Java on CPU(s) only or with @@ -166,7 +166,7 @@ Take the following steps to install TensorFlow for Java on Linux or macOS: OS=$(uname -s | tr '[:upper:]' '[:lower:]') mkdir -p ./jni curl -L \ - "https://storage.googleapis.com/tensorflow/libtensorflow/libtensorflow_jni-${TF_TYPE}-${OS}-x86_64-1.5.0-rc0.tar.gz" | + "https://storage.googleapis.com/tensorflow/libtensorflow/libtensorflow_jni-${TF_TYPE}-${OS}-x86_64-1.5.0-rc1.tar.gz" | tar -xz -C ./jni ### Install on Windows @@ -174,10 +174,10 @@ Take the following steps to install TensorFlow for Java on Linux or macOS: Take the following steps to install TensorFlow for Java on Windows: 1. Download - [libtensorflow.jar](https://storage.googleapis.com/tensorflow/libtensorflow/libtensorflow-1.5.0-rc0.jar), + [libtensorflow.jar](https://storage.googleapis.com/tensorflow/libtensorflow/libtensorflow-1.5.0-rc1.jar), which is the TensorFlow Java Archive (JAR). 2. Download the following Java Native Interface (JNI) file appropriate for - [TensorFlow for Java on Windows](https://storage.googleapis.com/tensorflow/libtensorflow/libtensorflow_jni-cpu-windows-x86_64-1.5.0-rc0.zip). + [TensorFlow for Java on Windows](https://storage.googleapis.com/tensorflow/libtensorflow/libtensorflow_jni-cpu-windows-x86_64-1.5.0-rc1.zip). 3. Extract this .zip file. @@ -225,7 +225,7 @@ must be part of your `classpath`. For example, you can include the downloaded `.jar` in your `classpath` by using the `-cp` compilation flag as follows: -<pre><b>javac -cp libtensorflow-1.5.0-rc0.jar HelloTF.java</b></pre> +<pre><b>javac -cp libtensorflow-1.5.0-rc1.jar HelloTF.java</b></pre> ### Running @@ -239,11 +239,11 @@ two files are available to the JVM: For example, the following command line executes the `HelloTF` program on Linux and macOS X: -<pre><b>java -cp libtensorflow-1.5.0-rc0.jar:. -Djava.library.path=./jni HelloTF</b></pre> +<pre><b>java -cp libtensorflow-1.5.0-rc1.jar:. -Djava.library.path=./jni HelloTF</b></pre> And the following command line executes the `HelloTF` program on Windows: -<pre><b>java -cp libtensorflow-1.5.0-rc0.jar;. -Djava.library.path=jni HelloTF</b></pre> +<pre><b>java -cp libtensorflow-1.5.0-rc1.jar;. -Djava.library.path=jni HelloTF</b></pre> If the program prints <tt>Hello from <i>version</i></tt>, you've successfully installed TensorFlow for Java and are ready to use the API. If the program diff --git a/tensorflow/docs_src/install/install_linux.md b/tensorflow/docs_src/install/install_linux.md index bb1d9a9f57..03f12dff08 100644 --- a/tensorflow/docs_src/install/install_linux.md +++ b/tensorflow/docs_src/install/install_linux.md @@ -188,7 +188,7 @@ Take the following steps to install TensorFlow with Virtualenv: Virtualenv environment: <pre>(tensorflow)$ <b>pip3 install --upgrade \ - https://storage.googleapis.com/tensorflow/linux/cpu/tensorflow-1.5.0rc0-cp34-cp34m-linux_x86_64.whl</b></pre> + https://storage.googleapis.com/tensorflow/linux/cpu/tensorflow-1.5.0rc1-cp34-cp34m-linux_x86_64.whl</b></pre> If you encounter installation problems, see [Common Installation Problems](#common_installation_problems). @@ -293,7 +293,7 @@ take the following steps: <pre> $ <b>sudo pip3 install --upgrade \ - https://storage.googleapis.com/tensorflow/linux/cpu/tensorflow-1.5.0rc0-cp34-cp34m-linux_x86_64.whl</b> + https://storage.googleapis.com/tensorflow/linux/cpu/tensorflow-1.5.0rc1-cp34-cp34m-linux_x86_64.whl</b> </pre> If this step fails, see @@ -480,7 +480,7 @@ Take the following steps to install TensorFlow in an Anaconda environment: <pre> (tensorflow)$ <b>pip install --ignore-installed --upgrade \ - https://storage.googleapis.com/tensorflow/linux/cpu/tensorflow-1.5.0rc0-cp34-cp34m-linux_x86_64.whl</b></pre> + https://storage.googleapis.com/tensorflow/linux/cpu/tensorflow-1.5.0rc1-cp34-cp34m-linux_x86_64.whl</b></pre> <a name="ValidateYourInstallation"></a> @@ -648,14 +648,14 @@ This section documents the relevant values for Linux installations. CPU only: <pre> -https://storage.googleapis.com/tensorflow/linux/cpu/tensorflow-1.5.0rc0-cp27-none-linux_x86_64.whl +https://storage.googleapis.com/tensorflow/linux/cpu/tensorflow-1.5.0rc1-cp27-none-linux_x86_64.whl </pre> GPU support: <pre> -https://storage.googleapis.com/tensorflow/linux/gpu/tensorflow_gpu-1.5.0rc0-cp27-none-linux_x86_64.whl +https://storage.googleapis.com/tensorflow/linux/gpu/tensorflow_gpu-1.5.0rc1-cp27-none-linux_x86_64.whl </pre> Note that GPU support requires the NVIDIA hardware and software described in @@ -667,14 +667,14 @@ Note that GPU support requires the NVIDIA hardware and software described in CPU only: <pre> -https://storage.googleapis.com/tensorflow/linux/cpu/tensorflow-1.5.0rc0-cp34-cp34m-linux_x86_64.whl +https://storage.googleapis.com/tensorflow/linux/cpu/tensorflow-1.5.0rc1-cp34-cp34m-linux_x86_64.whl </pre> GPU support: <pre> -https://storage.googleapis.com/tensorflow/linux/gpu/tensorflow_gpu-1.5.0rc0-cp34-cp34m-linux_x86_64.whl +https://storage.googleapis.com/tensorflow/linux/gpu/tensorflow_gpu-1.5.0rc1-cp34-cp34m-linux_x86_64.whl </pre> Note that GPU support requires the NVIDIA hardware and software described in @@ -686,14 +686,14 @@ Note that GPU support requires the NVIDIA hardware and software described in CPU only: <pre> -https://storage.googleapis.com/tensorflow/linux/cpu/tensorflow-1.5.0rc0-cp35-cp35m-linux_x86_64.whl +https://storage.googleapis.com/tensorflow/linux/cpu/tensorflow-1.5.0rc1-cp35-cp35m-linux_x86_64.whl </pre> GPU support: <pre> -https://storage.googleapis.com/tensorflow/linux/gpu/tensorflow_gpu-1.5.0rc0-cp35-cp35m-linux_x86_64.whl +https://storage.googleapis.com/tensorflow/linux/gpu/tensorflow_gpu-1.5.0rc1-cp35-cp35m-linux_x86_64.whl </pre> @@ -705,14 +705,14 @@ Note that GPU support requires the NVIDIA hardware and software described in CPU only: <pre> -https://storage.googleapis.com/tensorflow/linux/cpu/tensorflow-1.5.0rc0-cp36-cp36m-linux_x86_64.whl +https://storage.googleapis.com/tensorflow/linux/cpu/tensorflow-1.5.0rc1-cp36-cp36m-linux_x86_64.whl </pre> GPU support: <pre> -https://storage.googleapis.com/tensorflow/linux/gpu/tensorflow_gpu-1.5.0rc0-cp36-cp36m-linux_x86_64.whl +https://storage.googleapis.com/tensorflow/linux/gpu/tensorflow_gpu-1.5.0rc1-cp36-cp36m-linux_x86_64.whl </pre> diff --git a/tensorflow/docs_src/install/install_mac.md b/tensorflow/docs_src/install/install_mac.md index cf1c5157f8..e13ddadab7 100644 --- a/tensorflow/docs_src/install/install_mac.md +++ b/tensorflow/docs_src/install/install_mac.md @@ -115,7 +115,7 @@ Take the following steps to install TensorFlow with Virtualenv: TensorFlow in the active Virtualenv is as follows: <pre> $ <b>pip3 install --upgrade \ - https://storage.googleapis.com/tensorflow/mac/cpu/tensorflow-1.5.0rc0-py2-none-any.whl</b></pre> + https://storage.googleapis.com/tensorflow/mac/cpu/tensorflow-1.5.0rc1-py2-none-any.whl</b></pre> If you encounter installation problems, see [Common Installation Problems](#common-installation-problems). @@ -238,7 +238,7 @@ take the following steps: issue the following command: <pre> $ <b>sudo pip3 install --upgrade \ - https://storage.googleapis.com/tensorflow/mac/cpu/tensorflow-1.5.0rc0-py2-none-any.whl</b> </pre> + https://storage.googleapis.com/tensorflow/mac/cpu/tensorflow-1.5.0rc1-py2-none-any.whl</b> </pre> If the preceding command fails, see [installation problems](#common-installation-problems). @@ -347,7 +347,7 @@ Take the following steps to install TensorFlow in an Anaconda environment: TensorFlow for Python 2.7: <pre> (<i>targetDirectory</i>)$ <b>pip install --ignore-installed --upgrade \ - https://storage.googleapis.com/tensorflow/mac/cpu/tensorflow-1.5.0rc0-py2-none-any.whl</b></pre> + https://storage.googleapis.com/tensorflow/mac/cpu/tensorflow-1.5.0rc1-py2-none-any.whl</b></pre> <a name="ValidateYourInstallation"></a> @@ -520,7 +520,7 @@ This section documents the relevant values for Mac OS installations. <pre> -https://storage.googleapis.com/tensorflow/mac/cpu/tensorflow-1.5.0rc0-py2-none-any.whl +https://storage.googleapis.com/tensorflow/mac/cpu/tensorflow-1.5.0rc1-py2-none-any.whl </pre> @@ -528,5 +528,5 @@ https://storage.googleapis.com/tensorflow/mac/cpu/tensorflow-1.5.0rc0-py2-none-a <pre> -https://storage.googleapis.com/tensorflow/mac/cpu/tensorflow-1.5.0rc0-py3-none-any.whl +https://storage.googleapis.com/tensorflow/mac/cpu/tensorflow-1.5.0rc1-py3-none-any.whl </pre> diff --git a/tensorflow/docs_src/install/install_sources.md b/tensorflow/docs_src/install/install_sources.md index 90e93f56c5..f494cc7a7c 100644 --- a/tensorflow/docs_src/install/install_sources.md +++ b/tensorflow/docs_src/install/install_sources.md @@ -361,10 +361,10 @@ Invoke `pip install` to install that pip package. The filename of the `.whl` file depends on your platform. For example, the following command will install the pip package -for TensorFlow 1.5.0rc0 on Linux: +for TensorFlow 1.5.0rc1 on Linux: <pre> -$ <b>sudo pip install /tmp/tensorflow_pkg/tensorflow-1.5.0rc0-py2-none-any.whl</b> +$ <b>sudo pip install /tmp/tensorflow_pkg/tensorflow-1.5.0rc1-py2-none-any.whl</b> </pre> ## Validate your installation @@ -462,9 +462,12 @@ Stack Overflow and specify the `tensorflow` tag. **Linux** <table> <tr><th>Version:</th><th>CPU/GPU:</th><th>Python Version:</th><th>Compiler:</th><th>Build Tools:</th><th>cuDNN:</th><th>CUDA:</th></tr> + +<tr><td>tensorflow-1.5.0-rc1</td><td>CPU</td><td>2.7, 3.3-3.6</td><td>GCC 4.8</td><td>Bazel 0.8.0</td><td>N/A</td><td>N/A</td></tr> +<tr><td>tensorflow_gpu-1.5.0-rc1</td><td>GPU</td><td>2.7, 3.3-3.6</td><td>GCC 4.8</td><td>Bazel 0.8.0</td><td>7</td><td>9</td></tr> <tr><td>tensorflow-1.4.0</td><td>CPU</td><td>2.7, 3.3-3.6</td><td>GCC 4.8</td><td>Bazel 0.5.4</td><td>N/A</td><td>N/A</td></tr> <tr><td>tensorflow_gpu-1.4.0</td><td>GPU</td><td>2.7, 3.3-3.6</td><td>GCC 4.8</td><td>Bazel 0.5.4</td><td>6</td><td>8</td></tr> - <tr><td>tensorflow-1.3.0</td><td>CPU</td><td>2.7, 3.3-3.6</td><td>GCC 4.8</td><td>Bazel 0.4.5</td><td>N/A</td><td>N/A</td></tr> +<tr><td>tensorflow-1.3.0</td><td>CPU</td><td>2.7, 3.3-3.6</td><td>GCC 4.8</td><td>Bazel 0.4.5</td><td>N/A</td><td>N/A</td></tr> <tr><td>tensorflow_gpu-1.3.0</td><td>GPU</td><td>2.7, 3.3-3.6</td><td>GCC 4.8</td><td>Bazel 0.4.5</td><td>6</td><td>8</td></tr> <tr><td>tensorflow-1.2.0</td><td>CPU</td><td>2.7, 3.3-3.6</td><td>GCC 4.8</td><td>Bazel 0.4.5</td><td>N/A</td><td>N/A</td></tr> <tr><td>tensorflow_gpu-1.2.0</td><td>GPU</td><td>2.7, 3.3-3.6</td><td>GCC 4.8</td><td>Bazel 0.4.5</td><td>5.1</td><td>8</td></tr> @@ -477,8 +480,9 @@ Stack Overflow and specify the `tensorflow` tag. **Mac** <table> <tr><th>Version:</th><th>CPU/GPU:</th><th>Python Version:</th><th>Compiler:</th><th>Build Tools:</th><th>cuDNN:</th><th>CUDA:</th></tr> +<tr><td>tensorflow-1.5.0-rc1</td><td>CPU</td><td>2.7, 3.3-3.6</td><td>Clang from xcode</td><td>Bazel 0.8.1</td><td>N/A</td><td>N/A</td></tr> <tr><td>tensorflow-1.4.0</td><td>CPU</td><td>2.7, 3.3-3.6</td><td>Clang from xcode</td><td>Bazel 0.5.4</td><td>N/A</td><td>N/A</td></tr> - <tr><td>tensorflow-1.3.0</td><td>CPU</td><td>2.7, 3.3-3.6</td><td>Clang from xcode</td><td>Bazel 0.4.5</td><td>N/A</td><td>N/A</td></tr> +<tr><td>tensorflow-1.3.0</td><td>CPU</td><td>2.7, 3.3-3.6</td><td>Clang from xcode</td><td>Bazel 0.4.5</td><td>N/A</td><td>N/A</td></tr> <tr><td>tensorflow-1.2.0</td><td>CPU</td><td>2.7, 3.3-3.6</td><td>Clang from xcode</td><td>Bazel 0.4.5</td><td>N/A</td><td>N/A</td></tr> <tr><td>tensorflow-1.1.0</td><td>CPU</td><td>2.7, 3.3-3.6</td><td>Clang from xcode</td><td>Bazel 0.4.2</td><td>N/A</td><td>N/A</td></tr> <tr><td>tensorflow_gpu-1.1.0</td><td>GPU</td><td>2.7, 3.3-3.6</td><td>Clang from xcode</td><td>Bazel 0.4.2</td><td>5.1</td><td>8</td></tr> @@ -489,6 +493,8 @@ Stack Overflow and specify the `tensorflow` tag. **Windows** <table> <tr><th>Version:</th><th>CPU/GPU:</th><th>Python Version:</th><th>Compiler:</th><th>Build Tools:</th><th>cuDNN:</th><th>CUDA:</th></tr> +<tr><td>tensorflow-1.5.0-rc1</td><td>CPU</td><td>3.5-3.6</td><td>MSVC 2015 update 3</td><td>Cmake v3.6.3</td><td>N/A</td><td>N/A</td></tr> +<tr><td>tensorflow_gpu-1.5.0-rc1</td><td>GPU</td><td>3.5-3.6</td><td>MSVC 2015 update 3</td><td>Cmake v3.6.3</td><td>7</td><td>9</td></tr> <tr><td>tensorflow-1.4.0</td><td>CPU</td><td>3.5-3.6</td><td>MSVC 2015 update 3</td><td>Cmake v3.6.3</td><td>N/A</td><td>N/A</td></tr> <tr><td>tensorflow_gpu-1.4.0</td><td>GPU</td><td>3.5-3.6</td><td>MSVC 2015 update 3</td><td>Cmake v3.6.3</td><td>6</td><td>8</td></tr> <tr><td>tensorflow-1.3.0</td><td>CPU</td><td>3.5-3.6</td><td>MSVC 2015 update 3</td><td>Cmake v3.6.3</td><td>N/A</td><td>N/A</td></tr> diff --git a/tensorflow/examples/tutorials/word2vec/word2vec_basic.py b/tensorflow/examples/tutorials/word2vec/word2vec_basic.py index 87cd95165e..7d1650f05e 100644 --- a/tensorflow/examples/tutorials/word2vec/word2vec_basic.py +++ b/tensorflow/examples/tutorials/word2vec/word2vec_basic.py @@ -21,6 +21,8 @@ from __future__ import print_function import collections import math import os +import sys +import argparse import random from tempfile import gettempdir import zipfile @@ -30,6 +32,24 @@ from six.moves import urllib from six.moves import xrange # pylint: disable=redefined-builtin import tensorflow as tf +from tensorflow.contrib.tensorboard.plugins import projector + +# Give a folder path as an argument with '--log_dir' to save +# TensorBoard summaries. Default is a log folder in current directory. +current_path = os.path.dirname(os.path.realpath(sys.argv[0])) + +parser = argparse.ArgumentParser() +parser.add_argument( + '--log_dir', + type=str, + default=os.path.join(current_path, 'log'), + help='The log directory for TensorBoard summaries.') +FLAGS, unparsed = parser.parse_known_args() + +# Create the directory for TensorBoard variables if there is not. +if not os.path.exists(FLAGS.log_dir): + os.makedirs(FLAGS.log_dir) + # Step 1: Download the data. url = 'http://mattmahoney.net/dc/' @@ -156,38 +176,47 @@ graph = tf.Graph() with graph.as_default(): # Input data. - train_inputs = tf.placeholder(tf.int32, shape=[batch_size]) - train_labels = tf.placeholder(tf.int32, shape=[batch_size, 1]) - valid_dataset = tf.constant(valid_examples, dtype=tf.int32) + with tf.name_scope('inputs'): + train_inputs = tf.placeholder(tf.int32, shape=[batch_size]) + train_labels = tf.placeholder(tf.int32, shape=[batch_size, 1]) + valid_dataset = tf.constant(valid_examples, dtype=tf.int32) # Ops and variables pinned to the CPU because of missing GPU implementation with tf.device('/cpu:0'): # Look up embeddings for inputs. - embeddings = tf.Variable( - tf.random_uniform([vocabulary_size, embedding_size], -1.0, 1.0)) - embed = tf.nn.embedding_lookup(embeddings, train_inputs) + with tf.name_scope('embeddings'): + embeddings = tf.Variable( + tf.random_uniform([vocabulary_size, embedding_size], -1.0, 1.0)) + embed = tf.nn.embedding_lookup(embeddings, train_inputs) # Construct the variables for the NCE loss - nce_weights = tf.Variable( - tf.truncated_normal([vocabulary_size, embedding_size], - stddev=1.0 / math.sqrt(embedding_size))) - nce_biases = tf.Variable(tf.zeros([vocabulary_size])) + with tf.name_scope('weights'): + nce_weights = tf.Variable( + tf.truncated_normal([vocabulary_size, embedding_size], + stddev=1.0 / math.sqrt(embedding_size))) + with tf.name_scope('biases'): + nce_biases = tf.Variable(tf.zeros([vocabulary_size])) # Compute the average NCE loss for the batch. # tf.nce_loss automatically draws a new sample of the negative labels each # time we evaluate the loss. # Explanation of the meaning of NCE loss: # http://mccormickml.com/2016/04/19/word2vec-tutorial-the-skip-gram-model/ - loss = tf.reduce_mean( - tf.nn.nce_loss(weights=nce_weights, - biases=nce_biases, - labels=train_labels, - inputs=embed, - num_sampled=num_sampled, - num_classes=vocabulary_size)) + with tf.name_scope('loss'): + loss = tf.reduce_mean( + tf.nn.nce_loss(weights=nce_weights, + biases=nce_biases, + labels=train_labels, + inputs=embed, + num_sampled=num_sampled, + num_classes=vocabulary_size)) + + # Add the loss value as a scalar to summary. + tf.summary.scalar('loss', loss) # Construct the SGD optimizer using a learning rate of 1.0. - optimizer = tf.train.GradientDescentOptimizer(1.0).minimize(loss) + with tf.name_scope('optimizer'): + optimizer = tf.train.GradientDescentOptimizer(1.0).minimize(loss) # Compute the cosine similarity between minibatch examples and all embeddings. norm = tf.sqrt(tf.reduce_sum(tf.square(embeddings), 1, keep_dims=True)) @@ -197,13 +226,22 @@ with graph.as_default(): similarity = tf.matmul( valid_embeddings, normalized_embeddings, transpose_b=True) + # Merge all summaries. + merged = tf.summary.merge_all() + # Add variable initializer. init = tf.global_variables_initializer() + # Create a saver. + saver = tf.train.Saver() + # Step 5: Begin training. num_steps = 100001 with tf.Session(graph=graph) as session: + # Open a writer to write summaries. + writer = tf.summary.FileWriter(FLAGS.log_dir, session.graph) + # We must initialize all variables before we use them. init.run() print('Initialized') @@ -214,10 +252,21 @@ with tf.Session(graph=graph) as session: batch_size, num_skips, skip_window) feed_dict = {train_inputs: batch_inputs, train_labels: batch_labels} + # Define metadata variable. + run_metadata = tf.RunMetadata() + # We perform one update step by evaluating the optimizer op (including it # in the list of returned values for session.run() - _, loss_val = session.run([optimizer, loss], feed_dict=feed_dict) + # Also, evaluate the merged op to get all summaries from the returned "summary" variable. + # Feed metadata variable to session for visualizing the graph in TensorBoard. + _, summary, loss_val = session.run([optimizer, merged, loss], feed_dict=feed_dict, run_metadata=run_metadata) average_loss += loss_val + + # Add returned summaries to writer in each step. + writer.add_summary(summary, step) + # Add metadata to visualize the graph for the last run. + if step == (num_steps - 1): + writer.add_run_metadata(run_metadata, 'step%d' % step) if step % 2000 == 0: if step > 0: @@ -240,6 +289,23 @@ with tf.Session(graph=graph) as session: print(log_str) final_embeddings = normalized_embeddings.eval() + # Write corresponding labels for the embeddings. + with open(FLAGS.log_dir + '/metadata.tsv', 'w') as f: + for i in xrange(vocabulary_size): + f.write(reverse_dictionary[i] + '\n') + + # Save the model for checkpoints. + saver.save(session, os.path.join(FLAGS.log_dir, "model.ckpt")) + + # Create a configuration for visualizing embeddings with the labels in TensorBoard. + config = projector.ProjectorConfig() + embedding_conf = config.embeddings.add() + embedding_conf.tensor_name = embeddings.name + embedding_conf.metadata_path = os.path.join(FLAGS.log_dir, 'metadata.tsv') + projector.visualize_embeddings(writer, config) + +writer.close() + # Step 6: Visualize the embeddings. diff --git a/tensorflow/python/data/kernel_tests/batch_dataset_op_test.py b/tensorflow/python/data/kernel_tests/batch_dataset_op_test.py index 53c8be1d1d..eac1c1960d 100644 --- a/tensorflow/python/data/kernel_tests/batch_dataset_op_test.py +++ b/tensorflow/python/data/kernel_tests/batch_dataset_op_test.py @@ -1,3 +1,4 @@ +# -*- coding: utf-8 -*- # Copyright 2017 The TensorFlow Authors. All Rights Reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -301,6 +302,27 @@ class BatchDatasetTest(test.TestCase): with self.assertRaises(errors.OutOfRangeError): sess.run(get_next) + def testPaddedBatchDatasetUnicode(self): + # See GitHub issue 16149 + def generator(): + data = [ + [u'Простой', u'тест', u'юникода'], + [u'никогда', u'не', u'бывает', u'простым']] + + for seq in data: + yield seq, [0, 1, 2, 3] + + dataset = dataset_ops.Dataset.from_generator( + generator, + (dtypes.string, dtypes.int32), + (tensor_shape.TensorShape([None]), tensor_shape.TensorShape([None]))) + padded_dataset = dataset.padded_batch(2, padded_shapes=([None], [None]), + padding_values=('', 0)) + with self.test_session() as sess: + next_element = padded_dataset.make_one_shot_iterator().get_next() + sess.run(next_element) + + def testPaddedBatchDatasetShapeSpecifications(self): int_placeholder = array_ops.placeholder(dtypes.int32) float_placeholder = array_ops.placeholder(dtypes.float32) diff --git a/tensorflow/python/estimator/estimator.py b/tensorflow/python/estimator/estimator.py index b4c6b2747e..2e914fa7e0 100644 --- a/tensorflow/python/estimator/estimator.py +++ b/tensorflow/python/estimator/estimator.py @@ -128,9 +128,10 @@ class Estimator(object): model_dir: Directory to save model parameters, graph and etc. This can also be used to load checkpoints from the directory into a estimator to - continue training a previously saved model. If `None`, the model_dir in - `config` will be used if set. If both are set, they must be same. If - both are `None`, a temporary directory will be used. + continue training a previously saved model. If `PathLike` object, the + path will be resolved. If `None`, the model_dir in `config` will be used + if set. If both are set, they must be same. If both are `None`, a + temporary directory will be used. config: Configuration object. params: `dict` of hyper parameters that will be passed into `model_fn`. Keys are names of parameters, values are basic python types. @@ -158,6 +159,7 @@ class Estimator(object): self._config = config # Model directory. + model_dir = compat.path_to_str(model_dir) if (model_dir is not None) and (self._config.model_dir is not None): if model_dir != self._config.model_dir: # TODO(alanyee): remove this suppression after it is no longer needed diff --git a/tensorflow/python/estimator/run_config.py b/tensorflow/python/estimator/run_config.py index 894005550d..61a537022b 100644 --- a/tensorflow/python/estimator/run_config.py +++ b/tensorflow/python/estimator/run_config.py @@ -27,6 +27,7 @@ import six from tensorflow.core.protobuf import config_pb2 from tensorflow.python.platform import tf_logging as logging from tensorflow.python.training import server_lib +from tensorflow.python.util import compat _USE_DEFAULT = object() @@ -399,7 +400,8 @@ class RunConfig(object): Args: model_dir: directory where model parameters, graph, etc are saved. If - `None`, will use a default value set by the Estimator. + `PathLike` object, the path will be resolved. If `None`, will use a + default value set by the Estimator. tf_random_seed: Random seed for TensorFlow initializers. Setting this value allows consistency between reruns. save_summary_steps: Save summaries every this many steps. @@ -442,7 +444,7 @@ class RunConfig(object): if tf_config: logging.info('TF_CONFIG environment variable: %s', tf_config) - model_dir = _get_model_dir(tf_config, model_dir) + model_dir = _get_model_dir(tf_config, compat.path_to_str(model_dir)) RunConfig._replace( self, diff --git a/tensorflow/python/framework/constant_op.py b/tensorflow/python/framework/constant_op.py index 0a56d3f64d..d3d8c9c154 100644 --- a/tensorflow/python/framework/constant_op.py +++ b/tensorflow/python/framework/constant_op.py @@ -60,7 +60,6 @@ def _eager_reshape(tensor, shape, ctx): attr_t = tensor._datatype_enum() # pylint: disable=protected-access attr_tshape, (shape,) = execute.args_to_matching_eager( [shape], ctx, dtypes.int32) - attr_tshape = attr_tshape inputs_flat = [tensor, shape] attrs = ("T", attr_t, "Tshape", attr_tshape) result, = execute.execute( diff --git a/tensorflow/python/kernel_tests/conv1d_test.py b/tensorflow/python/kernel_tests/conv1d_test.py index d92797a7d3..e2e6205911 100644 --- a/tensorflow/python/kernel_tests/conv1d_test.py +++ b/tensorflow/python/kernel_tests/conv1d_test.py @@ -30,27 +30,29 @@ from tensorflow.python.platform import test class Conv1DTest(test.TestCase): def testBasic(self): - """Test that argument passing to conv2d is handled properly.""" - - x = constant_op.constant([1, 2, 3, 4], dtype=dtypes.float32) - x = array_ops.expand_dims(x, 0) # Add batch dimension - x = array_ops.expand_dims(x, 2) # And depth dimension - filters = constant_op.constant([2, 1], dtype=dtypes.float32) - filters = array_ops.expand_dims(filters, 1) # in_channels - filters = array_ops.expand_dims(filters, 2) # out_channels - # Filters is 2x1x1 - for stride in [1, 2]: - with self.test_session(use_gpu=test.is_gpu_available()): - c = nn_ops.conv1d(x, filters, stride, padding="VALID") - reduced = array_ops.squeeze(c) - output = reduced.eval() - if stride == 1: - self.assertEqual(len(output), 3) - self.assertAllClose(output, - [2 * 1 + 1 * 2, 2 * 2 + 1 * 3, 2 * 3 + 1 * 4]) - else: - self.assertEqual(len(output), 2) - self.assertAllClose(output, [2 * 1 + 1 * 2, 2 * 3 + 1 * 4]) + """Test that argument passing to conv1d is handled properly.""" + # TODO(yongtang): dtypes.float64 can only be enabled once conv2d support + # dtypes.float64, as conv1d implicitly calls conv2d after expand_dims. + for dtype in [dtypes.float16, dtypes.float32]: + x = constant_op.constant([1, 2, 3, 4], dtype=dtype) + x = array_ops.expand_dims(x, 0) # Add batch dimension + x = array_ops.expand_dims(x, 2) # And depth dimension + filters = constant_op.constant([2, 1], dtype=dtype) + filters = array_ops.expand_dims(filters, 1) # in_channels + filters = array_ops.expand_dims(filters, 2) # out_channels + # Filters is 2x1x1 + for stride in [1, 2]: + with self.test_session(use_gpu=test.is_gpu_available()): + c = nn_ops.conv1d(x, filters, stride, padding="VALID") + reduced = array_ops.squeeze(c) + output = reduced.eval() + if stride == 1: + self.assertEqual(len(output), 3) + self.assertAllClose(output, + [2 * 1 + 1 * 2, 2 * 2 + 1 * 3, 2 * 3 + 1 * 4]) + else: + self.assertEqual(len(output), 2) + self.assertAllClose(output, [2 * 1 + 1 * 2, 2 * 3 + 1 * 4]) def testConv1DTranspose(self): with self.test_session(): diff --git a/tensorflow/python/kernel_tests/cwise_ops_test.py b/tensorflow/python/kernel_tests/cwise_ops_test.py index cea12ea8ec..a91917b27f 100644 --- a/tensorflow/python/kernel_tests/cwise_ops_test.py +++ b/tensorflow/python/kernel_tests/cwise_ops_test.py @@ -24,6 +24,7 @@ import numpy as np from tensorflow.python.framework import constant_op from tensorflow.python.framework import dtypes as dtypes_lib +from tensorflow.python.framework import errors_impl from tensorflow.python.framework import ops from tensorflow.python.framework import sparse_tensor from tensorflow.python.ops import array_ops @@ -1168,6 +1169,32 @@ class BinaryOpTest(test.TestCase): self._compareCpu(x1, x2, np.arctan2, math_ops.atan2) self._compareGpu(x1, x2, np.arctan2, math_ops.atan2) + def testPowNegativeExponent(self): + for dtype in [np.int32, np.int64]: + with self.test_session(use_gpu=False) as sess: + with self.assertRaisesRegexp( + errors_impl.InvalidArgumentError, + "Integers to negative integer powers are not allowed"): + x = np.array([5, 2]).astype(dtype) + y = np.array([-2, 3]).astype(dtype) + sess.run(math_ops.pow(x, y)) + + with self.test_session(use_gpu=False) as sess: + with self.assertRaisesRegexp( + errors_impl.InvalidArgumentError, + "Integers to negative integer powers are not allowed"): + x = np.array([5, 2]).astype(dtype) + y = np.array([2, -3]).astype(dtype) + sess.run(math_ops.pow(x, y)) + + with self.test_session(use_gpu=False) as sess: + with self.assertRaisesRegexp( + errors_impl.InvalidArgumentError, + "Integers to negative integer powers are not allowed"): + x = np.array([5, 2]).astype(dtype) + y = -3 + sess.run(math_ops.pow(x, y)) + class ComparisonOpTest(test.TestCase): diff --git a/tensorflow/python/kernel_tests/metrics_test.py b/tensorflow/python/kernel_tests/metrics_test.py index 3358b78efd..e0e752147c 100644 --- a/tensorflow/python/kernel_tests/metrics_test.py +++ b/tensorflow/python/kernel_tests/metrics_test.py @@ -3628,7 +3628,8 @@ class MeanPerClassAccuracyTest(test.TestCase): predictions=array_ops.ones([10, 1]), labels=array_ops.ones([10, 1]), num_classes=2) - _assert_metric_variables(self, ('mean_accuracy/total_confusion_matrix:0',)) + _assert_metric_variables(self, ('mean_accuracy/count:0', + 'mean_accuracy/total:0')) def testMetricsCollections(self): my_collection_name = '__metrics__' @@ -3797,23 +3798,6 @@ class MeanPerClassAccuracyTest(test.TestCase): desired_output = np.mean([1.0 / 2.0, 2.0 / 3.0, 0.]) self.assertAlmostEqual(desired_output, mean_accuracy.eval()) - def testUpdateOpEvalIsAccumulatedConfusionMatrix(self): - predictions = array_ops.concat([ - constant_op.constant(0, shape=[5]), constant_op.constant(1, shape=[5]) - ], 0) - labels = array_ops.concat([ - constant_op.constant(0, shape=[3]), constant_op.constant(1, shape=[7]) - ], 0) - num_classes = 2 - with self.test_session() as sess: - mean_accuracy, update_op = metrics.mean_per_class_accuracy( - labels, predictions, num_classes) - sess.run(variables.local_variables_initializer()) - confusion_matrix = update_op.eval() - self.assertAllEqual([[3, 0], [2, 5]], confusion_matrix) - desired_mean_accuracy = np.mean([3. / 3., 5. / 7.]) - self.assertAlmostEqual(desired_mean_accuracy, mean_accuracy.eval()) - def testAllCorrect(self): predictions = array_ops.zeros([40]) labels = array_ops.zeros([40]) @@ -3822,7 +3806,7 @@ class MeanPerClassAccuracyTest(test.TestCase): mean_accuracy, update_op = metrics.mean_per_class_accuracy( labels, predictions, num_classes) sess.run(variables.local_variables_initializer()) - self.assertEqual(40, update_op.eval()[0]) + self.assertEqual(1.0, update_op.eval()[0]) self.assertEqual(1.0, mean_accuracy.eval()) def testAllWrong(self): @@ -3833,7 +3817,7 @@ class MeanPerClassAccuracyTest(test.TestCase): mean_accuracy, update_op = metrics.mean_per_class_accuracy( labels, predictions, num_classes) sess.run(variables.local_variables_initializer()) - self.assertAllEqual([[0, 0], [40, 0]], update_op.eval()) + self.assertAllEqual([0.0, 0.0], update_op.eval()) self.assertEqual(0., mean_accuracy.eval()) def testResultsWithSomeMissing(self): @@ -3852,8 +3836,9 @@ class MeanPerClassAccuracyTest(test.TestCase): mean_accuracy, update_op = metrics.mean_per_class_accuracy( labels, predictions, num_classes, weights=weights) sess.run(variables.local_variables_initializer()) - self.assertAllEqual([[2, 0], [2, 4]], update_op.eval()) - desired_mean_accuracy = np.mean([2. / 2., 4. / 6.]) + desired_accuracy = np.array([2. / 2., 4. / 6.], dtype=np.float32) + self.assertAllEqual(desired_accuracy, update_op.eval()) + desired_mean_accuracy = np.mean(desired_accuracy) self.assertAlmostEqual(desired_mean_accuracy, mean_accuracy.eval()) diff --git a/tensorflow/python/kernel_tests/xent_op_test.py b/tensorflow/python/kernel_tests/xent_op_test.py index 43be08f8a1..c6c7c4e26c 100644 --- a/tensorflow/python/kernel_tests/xent_op_test.py +++ b/tensorflow/python/kernel_tests/xent_op_test.py @@ -240,6 +240,16 @@ class XentTest(test.TestCase): self._testXentWrapper(features, labels, dim=-1, use_gpu=False) self._testXentWrapper(features, labels, dim=-1, use_gpu=True) + def testZeroDimension(self): + features = np.zeros([0, 2, 4]).astype(np.float32) + labels = np.zeros([0, 2, 4]).astype(np.float32) + np_loss, _ = self._npXent(features, labels) + with self.test_session(use_gpu=True) as sess: + loss = nn_ops.softmax_cross_entropy_with_logits( + labels=labels, logits=features) + tf_loss = sess.run(loss) + self.assertAllEqual(np_loss, tf_loss) + if __name__ == "__main__": test.main() diff --git a/tensorflow/python/layers/pooling.py b/tensorflow/python/layers/pooling.py index c6bd7aae07..ab06a3a408 100644 --- a/tensorflow/python/layers/pooling.py +++ b/tensorflow/python/layers/pooling.py @@ -63,14 +63,18 @@ class _Pooling1D(base.Layer): def call(self, inputs): # There is no TF op for 1D pooling, hence we make the inputs 4D. if self.data_format == 'channels_last': - inputs = array_ops.expand_dims(inputs, 2) - pool_shape = (1,) + self.pool_size + (1, 1) - strides = (1,) + self.strides + (1, 1) - data_format = 'NHWC' - else: + # input is NWC, make it NHWC inputs = array_ops.expand_dims(inputs, 1) + # pool on the W dim pool_shape = (1, 1) + self.pool_size + (1,) strides = (1, 1) + self.strides + (1,) + data_format = 'NHWC' + else: + # input is NCW, make it NCHW + inputs = array_ops.expand_dims(inputs, 2) + # pool on the W dim + pool_shape = (1, 1, 1) + self.pool_size + strides = (1, 1, 1) + self.strides data_format = 'NCHW' outputs = self.pool_function( @@ -81,9 +85,9 @@ class _Pooling1D(base.Layer): data_format=data_format) if self.data_format == 'channels_last': - return array_ops.squeeze(outputs, 2) - else: return array_ops.squeeze(outputs, 1) + else: + return array_ops.squeeze(outputs, 2) def compute_output_shape(self, input_shape): input_shape = tensor_shape.TensorShape(input_shape).as_list() diff --git a/tensorflow/python/layers/pooling_test.py b/tensorflow/python/layers/pooling_test.py index 589fee5f71..e4d4ed4a2a 100644 --- a/tensorflow/python/layers/pooling_test.py +++ b/tensorflow/python/layers/pooling_test.py @@ -110,19 +110,19 @@ class PoolingTest(test.TestCase): def testCreateMaxPooling1DChannelsFirst(self): width = 7 - images = random_ops.random_uniform((5, width, 4)) + images = random_ops.random_uniform((5, 4, width)) layer = pooling_layers.MaxPooling1D( 2, strides=2, data_format='channels_first') output = layer.apply(images) - self.assertListEqual(output.get_shape().as_list(), [5, 3, 4]) + self.assertListEqual(output.get_shape().as_list(), [5, 4, 3]) def testCreateAveragePooling1DChannelsFirst(self): width = 7 - images = random_ops.random_uniform((5, width, 4)) + images = random_ops.random_uniform((5, 4, width)) layer = pooling_layers.AveragePooling1D( 2, strides=2, data_format='channels_first') output = layer.apply(images) - self.assertListEqual(output.get_shape().as_list(), [5, 3, 4]) + self.assertListEqual(output.get_shape().as_list(), [5, 4, 3]) def testCreateMaxPooling3D(self): depth, height, width = 6, 7, 9 diff --git a/tensorflow/python/lib/core/py_func.cc b/tensorflow/python/lib/core/py_func.cc index dc56b39486..d3bfa0ee33 100644 --- a/tensorflow/python/lib/core/py_func.cc +++ b/tensorflow/python/lib/core/py_func.cc @@ -31,6 +31,7 @@ limitations under the License. #include "tensorflow/python/lib/core/ndarray_tensor_bridge.h" #include "tensorflow/python/lib/core/py_util.h" #include "tensorflow/python/lib/core/safe_ptr.h" + #include <Python.h> namespace tensorflow { @@ -141,7 +142,8 @@ bool IsSingleNone(PyObject* obj) { return false; } std::array<npy_intp, 0> indices; - char* item_ptr = static_cast<char*>(PyArray_GetPtr(array_obj, indices.data())); + char* item_ptr = + static_cast<char*>(PyArray_GetPtr(array_obj, indices.data())); PyObject* item = PyArray_GETITEM(array_obj, item_ptr); CHECK(item); return item == Py_None; @@ -301,13 +303,22 @@ Status ConvertNdarrayToTensor(PyObject* obj, Tensor* ret) { if (PyBytes_AsStringAndSize(input_data[i], &el, &el_size) == -1) { #if PY_MAJOR_VERSION >= 3 el = PyUnicode_AsUTF8AndSize(input_data[i], &el_size); - if (!el) { +#else + el = nullptr; + if (PyUnicode_Check(input_data[i])) { + PyObject* unicode = PyUnicode_AsUTF8String(input_data[i]); + if (unicode) { + if (PyString_AsStringAndSize(unicode, &el, &el_size) == -1) { + Py_DECREF(unicode); + el = nullptr; + } + } + } #endif + if (!el) { return errors::Unimplemented("Unsupported object type ", input_data[i]->ob_type->tp_name); -#if PY_MAJOR_VERSION >= 3 } -#endif } tflat(i) = string(el, el_size); } diff --git a/tensorflow/python/ops/histogram_ops.py b/tensorflow/python/ops/histogram_ops.py index d46084a41f..b2de2e5015 100644 --- a/tensorflow/python/ops/histogram_ops.py +++ b/tensorflow/python/ops/histogram_ops.py @@ -17,6 +17,7 @@ Please see @{$python/histogram_ops} guide. +@@histogram_fixed_width_bins @@histogram_fixed_width """ @@ -33,6 +34,70 @@ from tensorflow.python.ops import math_ops from tensorflow.python.util.tf_export import tf_export +def histogram_fixed_width_bins(values, + value_range, + nbins=100, + dtype=dtypes.int32, + name=None): + """Bins the given values for use in a histogram. + + Given the tensor `values`, this operation returns a rank 1 `Tensor` + representing the indices of a histogram into which each element + of `values` would be binned. The bins are equal width and + determined by the arguments `value_range` and `nbins`. + + Args: + values: Numeric `Tensor`. + value_range: Shape [2] `Tensor` of same `dtype` as `values`. + values <= value_range[0] will be mapped to hist[0], + values >= value_range[1] will be mapped to hist[-1]. + nbins: Scalar `int32 Tensor`. Number of histogram bins. + dtype: dtype for returned histogram. + name: A name for this operation (defaults to 'histogram_fixed_width'). + + Returns: + A `Tensor` holding the indices of the binned values whose shape matches + `values`. + + Examples: + + ```python + # Bins will be: (-inf, 1), [1, 2), [2, 3), [3, 4), [4, inf) + nbins = 5 + value_range = [0.0, 5.0] + new_values = [-1.0, 0.0, 1.5, 2.0, 5.0, 15] + + with tf.get_default_session() as sess: + indices = tf.histogram_fixed_width_bins(new_values, value_range, nbins=5) + variables.global_variables_initializer().run() + sess.run(indices) => [0, 0, 1, 2, 4] + ``` + """ + with ops.name_scope(name, 'histogram_fixed_width_bins', + [values, value_range, nbins]) as scope: + values = ops.convert_to_tensor(values, name='values') + shape = array_ops.shape(values) + + values = array_ops.reshape(values, [-1]) + value_range = ops.convert_to_tensor(value_range, name='value_range') + nbins = ops.convert_to_tensor(nbins, dtype=dtypes.int32, name='nbins') + nbins_float = math_ops.cast(nbins, values.dtype) + + # Map tensor values that fall within value_range to [0, 1]. + scaled_values = math_ops.truediv(values - value_range[0], + value_range[1] - value_range[0], + name='scaled_values') + + # map tensor values within the open interval value_range to {0,.., nbins-1}, + # values outside the open interval will be zero or less, or nbins or more. + indices = math_ops.floor(nbins_float * scaled_values, name='indices') + + # Clip edge cases (e.g. value = value_range[1]) or "outliers." + indices = math_ops.cast( + clip_ops.clip_by_value(indices, 0, nbins_float - 1), dtypes.int32) + return array_ops.reshape(indices, shape) + + @tf_export('histogram_fixed_width') def histogram_fixed_width(values, value_range, diff --git a/tensorflow/python/ops/histogram_ops_test.py b/tensorflow/python/ops/histogram_ops_test.py index 19ad6cd2ba..80ee090575 100644 --- a/tensorflow/python/ops/histogram_ops_test.py +++ b/tensorflow/python/ops/histogram_ops_test.py @@ -21,11 +21,64 @@ from __future__ import print_function import numpy as np from tensorflow.python.framework import dtypes +from tensorflow.python.framework import constant_op from tensorflow.python.ops import array_ops from tensorflow.python.ops import histogram_ops from tensorflow.python.platform import test +class BinValuesFixedWidth(test.TestCase): + + def test_empty_input_gives_all_zero_counts(self): + # Bins will be: + # (-inf, 1), [1, 2), [2, 3), [3, 4), [4, inf) + value_range = [0.0, 5.0] + values = [] + expected_bins = [] + with self.test_session(): + bins = histogram_ops.histogram_fixed_width_bins(values, value_range, nbins=5) + self.assertEqual(dtypes.int32, bins.dtype) + self.assertAllClose(expected_bins, bins.eval()) + + def test_1d_values_int32_output(self): + # Bins will be: + # (-inf, 1), [1, 2), [2, 3), [3, 4), [4, inf) + value_range = [0.0, 5.0] + values = [-1.0, 0.0, 1.5, 2.0, 5.0, 15] + expected_bins = [0, 0, 1, 2, 4, 4] + with self.test_session(): + bins = histogram_ops.histogram_fixed_width_bins( + values, value_range, nbins=5, dtype=dtypes.int64) + self.assertEqual(dtypes.int32, bins.dtype) + self.assertAllClose(expected_bins, bins.eval()) + + def test_1d_float64_values_int32_output(self): + # Bins will be: + # (-inf, 1), [1, 2), [2, 3), [3, 4), [4, inf) + value_range = np.float64([0.0, 5.0]) + values = np.float64([-1.0, 0.0, 1.5, 2.0, 5.0, 15]) + expected_bins = [0, 0, 1, 2, 4, 4] + with self.test_session(): + bins = histogram_ops.histogram_fixed_width_bins( + values, value_range, nbins=5) + self.assertEqual(dtypes.int32, bins.dtype) + self.assertAllClose(expected_bins, bins.eval()) + + def test_2d_values(self): + # Bins will be: + # (-inf, 1), [1, 2), [2, 3), [3, 4), [4, inf) + value_range = [0.0, 5.0] + values = constant_op.constant( + [[-1.0, 0.0, 1.5], [2.0, 5.0, 15]], + shape=(2, 3)) + expected_bins = [[0, 0, 1], [2, 4, 4]] + with self.test_session(): + bins = histogram_ops.histogram_fixed_width_bins( + values, value_range, nbins=5) + self.assertEqual(dtypes.int32, bins.dtype) + self.assertAllClose(expected_bins, bins.eval()) + + class HistogramFixedWidthTest(test.TestCase): def setUp(self): diff --git a/tensorflow/python/ops/image_ops_impl.py b/tensorflow/python/ops/image_ops_impl.py index d860a3b618..b713c44717 100644 --- a/tensorflow/python/ops/image_ops_impl.py +++ b/tensorflow/python/ops/image_ops_impl.py @@ -148,6 +148,28 @@ def _Check3DImage(image, require_static=True): return [] +def _Assert3DImage(image): + """Assert that we are working with a properly shaped image. + + Performs the check statically if possible (i.e. if the shape + is statically known). Otherwise adds a control dependency + to an assert op that checks the dynamic shape. + + Args: + image: 3-D Tensor of shape [height, width, channels] + + Raises: + ValueError: if `image.shape` is not a 3-vector. + + Returns: + If the shape of `image` could be verified statically, `image` is + returned unchanged, otherwise there will be a control dependency + added that asserts the correct dynamic shape. + """ + return control_flow_ops.with_dependencies( + _Check3DImage(image, require_static=False), image) + + def _CheckAtLeast3DImage(image, require_static=True): """Assert that we are working with properly shaped image. @@ -223,8 +245,7 @@ def random_flip_up_down(image, seed=None): """ with ops.name_scope(None, 'random_flip_up_down', [image]) as scope: image = ops.convert_to_tensor(image, name='image') - image = control_flow_ops.with_dependencies( - _Check3DImage(image, require_static=False), image) + image = _Assert3DImage(image) uniform_random = random_ops.random_uniform([], 0, 1.0, seed=seed) mirror_cond = math_ops.less(uniform_random, .5) result = control_flow_ops.cond(mirror_cond, @@ -255,8 +276,7 @@ def random_flip_left_right(image, seed=None): """ with ops.name_scope(None, 'random_flip_left_right', [image]) as scope: image = ops.convert_to_tensor(image, name='image') - image = control_flow_ops.with_dependencies( - _Check3DImage(image, require_static=False), image) + image = _Assert3DImage(image) uniform_random = random_ops.random_uniform([], 0, 1.0, seed=seed) mirror_cond = math_ops.less(uniform_random, .5) result = control_flow_ops.cond(mirror_cond, @@ -286,8 +306,7 @@ def flip_left_right(image): """ with ops.name_scope(None, 'flip_left_right', [image]) as scope: image = ops.convert_to_tensor(image, name='image') - image = control_flow_ops.with_dependencies( - _Check3DImage(image, require_static=False), image) + image = _Assert3DImage(image) return fix_image_flip_shape(image, array_ops.reverse(image, [1], name=scope)) @@ -312,8 +331,7 @@ def flip_up_down(image): """ with ops.name_scope(None, 'flip_up_down', [image]) as scope: image = ops.convert_to_tensor(image, name='image') - image = control_flow_ops.with_dependencies( - _Check3DImage(image, require_static=False), image) + image = _Assert3DImage(image) return fix_image_flip_shape(image, array_ops.reverse(image, [0], name=scope)) @@ -332,8 +350,7 @@ def rot90(image, k=1, name=None): """ with ops.name_scope(name, 'rot90', [image, k]) as scope: image = ops.convert_to_tensor(image, name='image') - image = control_flow_ops.with_dependencies( - _Check3DImage(image, require_static=False), image) + image = _Assert3DImage(image) k = ops.convert_to_tensor(k, dtype=dtypes.int32, name='k') k.get_shape().assert_has_rank(0) k = math_ops.mod(k, 4) @@ -373,8 +390,7 @@ def transpose_image(image): """ with ops.name_scope(None, 'transpose_image', [image]) as scope: image = ops.convert_to_tensor(image, name='image') - image = control_flow_ops.with_dependencies( - _Check3DImage(image, require_static=False), image) + image = _Assert3DImage(image) return array_ops.transpose(image, [1, 0, 2], name=scope) @@ -410,8 +426,7 @@ def central_crop(image, central_fraction): if central_fraction == 1.0: return image - image = control_flow_ops.with_dependencies( - _Check3DImage(image, require_static=False), image) + image = _Assert3DImage(image) img_shape = array_ops.shape(image) depth = image.get_shape()[2] @@ -848,8 +863,7 @@ def per_image_standardization(image): """ with ops.name_scope(None, 'per_image_standardization', [image]) as scope: image = ops.convert_to_tensor(image, name='image') - image = control_flow_ops.with_dependencies( - _Check3DImage(image, require_static=False), image) + image = _Assert3DImage(image) num_pixels = math_ops.reduce_prod(array_ops.shape(image)) image = math_ops.cast(image, dtype=dtypes.float32) diff --git a/tensorflow/python/ops/image_ops_test.py b/tensorflow/python/ops/image_ops_test.py index 3a49d41c9e..80911ffe07 100644 --- a/tensorflow/python/ops/image_ops_test.py +++ b/tensorflow/python/ops/image_ops_test.py @@ -2833,6 +2833,16 @@ class PngTest(test_util.TensorFlowTestCase): class GifTest(test_util.TensorFlowTestCase): + def testOptimizedGifErrorString(self): + filename = "tensorflow/core/lib/gif/testdata/optimized.gif" + + with self.test_session(use_gpu=True) as sess: + gif = io_ops.read_file(filename) + image = image_ops.decode_gif(gif) + with self.assertRaisesRegexp( + errors.InvalidArgumentError, "can't process optimized gif"): + gif, image = sess.run([gif, image]) + def testValid(self): # Read some real GIFs prefix = "tensorflow/core/lib/gif/testdata/" diff --git a/tensorflow/python/ops/losses/losses_impl.py b/tensorflow/python/ops/losses/losses_impl.py index e292bccc2c..72508eb435 100644 --- a/tensorflow/python/ops/losses/losses_impl.py +++ b/tensorflow/python/ops/losses/losses_impl.py @@ -512,7 +512,7 @@ def mean_pairwise_squared_error( Raises: ValueError: If the shape of `predictions` doesn't match that of `labels` or - if the shape of `weights` is invalid. Also if `labels` or `predictions + if the shape of `weights` is invalid. Also if `labels` or `predictions` is None. """ if labels is None: diff --git a/tensorflow/python/ops/metrics_impl.py b/tensorflow/python/ops/metrics_impl.py index 92b3ff2250..2d77e26081 100644 --- a/tensorflow/python/ops/metrics_impl.py +++ b/tensorflow/python/ops/metrics_impl.py @@ -831,8 +831,8 @@ def mean_per_class_accuracy(labels, Calculates the accuracy for each class, then takes the mean of that. For estimation of the metric over a stream of data, the function creates an - `update_op` operation that updates these variables and returns the - `mean_accuracy`. + `update_op` operation that updates the accuracy of each class and returns + them. If `weights` is `None`, weights default to 1. Use weights of 0 to mask values. @@ -843,8 +843,8 @@ def mean_per_class_accuracy(labels, shape is [batch size] and type `int32` or `int64`. The tensor will be flattened if its rank > 1. num_classes: The possible number of labels the prediction task can - have. This value must be provided, since a confusion matrix of - dimension = [num_classes, num_classes] will be allocated. + have. This value must be provided, since two variables with shape = + [num_classes] will be allocated. weights: Optional `Tensor` whose rank is either 0, or the same rank as `labels`, and must be broadcastable to `labels` (i.e., all dimensions must be either `1`, or the same as the corresponding `labels` dimension). @@ -857,7 +857,7 @@ def mean_per_class_accuracy(labels, Returns: mean_accuracy: A `Tensor` representing the mean per class accuracy. - update_op: An operation that increments the confusion matrix. + update_op: An operation that updates the accuracy tensor. Raises: ValueError: If `predictions` and `labels` have mismatched shapes, or if @@ -872,27 +872,43 @@ def mean_per_class_accuracy(labels, with variable_scope.variable_scope(name, 'mean_accuracy', (predictions, labels, weights)): + labels = math_ops.to_int64(labels) + + # Flatten the input if its rank > 1. + if labels.get_shape().ndims > 1: + labels = array_ops.reshape(labels, [-1]) + + if predictions.get_shape().ndims > 1: + predictions = array_ops.reshape(predictions, [-1]) + # Check if shape is compatible. predictions.get_shape().assert_is_compatible_with(labels.get_shape()) - total_cm, update_op = _streaming_confusion_matrix( - labels, predictions, num_classes, weights=weights) + total = metric_variable([num_classes], dtypes.float32, name='total') + count = metric_variable([num_classes], dtypes.float32, name='count') - def compute_mean_accuracy(name): - """Compute the mean per class accuracy via the confusion matrix.""" - per_row_sum = math_ops.to_float(math_ops.reduce_sum(total_cm, 1)) - cm_diag = math_ops.to_float(array_ops.diag_part(total_cm)) - denominator = per_row_sum + ones = array_ops.ones([array_ops.size(labels)], dtypes.float32) - # If the value of the denominator is 0, set it to 1 to avoid - # zero division. - denominator = array_ops.where( - math_ops.greater(denominator, 0), denominator, - array_ops.ones_like(denominator)) - accuracies = math_ops.div(cm_diag, denominator) - return math_ops.reduce_mean(accuracies, name=name) + if labels.dtype != predictions.dtype: + predictions = math_ops.cast(predictions, labels.dtype) + is_correct = math_ops.to_float(math_ops.equal(predictions, labels)) + + if weights is not None: + if weights.get_shape().ndims > 1: + weights = array_ops.reshape(weights, [-1]) + weights = math_ops.to_float(weights) + + is_correct = is_correct * weights + ones = ones * weights + + update_total_op = state_ops.scatter_add(total, labels, ones) + update_count_op = state_ops.scatter_add(count, labels, is_correct) + + per_class_accuracy = _safe_div(count, total, None) - mean_accuracy_v = compute_mean_accuracy('mean_accuracy') + mean_accuracy_v = math_ops.reduce_mean(per_class_accuracy, + name='mean_accuracy') + update_op = _safe_div(update_count_op, update_total_op, name='update_op') if metrics_collections: ops.add_to_collections(metrics_collections, mean_accuracy_v) diff --git a/tensorflow/python/ops/nn_impl.py b/tensorflow/python/ops/nn_impl.py index 67eee1c29e..837ee02e64 100644 --- a/tensorflow/python/ops/nn_impl.py +++ b/tensorflow/python/ops/nn_impl.py @@ -196,7 +196,10 @@ def weighted_cross_entropy_with_logits(targets, logits, pos_weight, name=None): targets * -log(sigmoid(logits)) + (1 - targets) * -log(1 - sigmoid(logits)) - The argument `pos_weight` is used as a multiplier for the positive targets: + A value `pos_weights > 1` decreases the false negative count, hence increasing the recall. + Conversely setting `pos_weights < 1` decreases the false positive count and increases the precision. + This can be seen from the fact that `pos_weight` is introduced as a multiplicative coefficient for the positive targets term + in the loss expression: targets * -log(sigmoid(logits)) * pos_weight + (1 - targets) * -log(1 - sigmoid(logits)) diff --git a/tensorflow/python/ops/nn_ops.py b/tensorflow/python/ops/nn_ops.py index 09aa45dae1..32b14f86b5 100644 --- a/tensorflow/python/ops/nn_ops.py +++ b/tensorflow/python/ops/nn_ops.py @@ -1558,7 +1558,8 @@ def leaky_relu(features, alpha=0.2, name=None): http://web.stanford.edu/~awni/papers/relu_hybrid_icml2013_final.pdf Args: - features: A `Tensor` representing preactivation values. + features: A `Tensor` representing preactivation values. Must be one of + the following types: `float16`, `float32`, `float64`, `int32`, `int64`. alpha: Slope of the activation function at x < 0. name: A name for the operation (optional). @@ -1567,7 +1568,9 @@ def leaky_relu(features, alpha=0.2, name=None): """ with ops.name_scope(name, "LeakyRelu", [features, alpha]): features = ops.convert_to_tensor(features, name="features") - alpha = ops.convert_to_tensor(alpha, name="alpha") + if features.dtype.is_integer: + features = math_ops.to_float(features) + alpha = ops.convert_to_tensor(alpha, dtype=features.dtype, name="alpha") return math_ops.maximum(alpha * features, features) @@ -2323,7 +2326,7 @@ def conv1d(value, filters, stride, padding, returned to the caller. Args: - value: A 3D `Tensor`. Must be of type `float32` or `float64`. + value: A 3D `Tensor`. Must be of type `float16` or `float32`. filters: A 3D `Tensor`. Must have the same type as `input`. stride: An `integer`. The number of entries by which the filter is moved right at each step. diff --git a/tensorflow/python/ops/nn_test.py b/tensorflow/python/ops/nn_test.py index 66bc0803b7..6767564024 100644 --- a/tensorflow/python/ops/nn_test.py +++ b/tensorflow/python/ops/nn_test.py @@ -878,11 +878,13 @@ class LeakyReluTest(test_lib.TestCase): self.assertAllClose(inputs, outputs) def testValues(self): - np_values = np.array([-1.0, 0.0, 0.5, 1.0, 2.0], dtype=np.float32) - outputs = nn_ops.leaky_relu(constant_op.constant(np_values)) - with self.test_session() as sess: - outputs = sess.run(outputs) - self.assertAllClose(outputs, [-0.2, 0.0, 0.5, 1.0, 2.0]) + for dtype in [np.int32, np.int64, np.float16, np.float32, np.float64]: + np_values = np.array([-2, -1, 0, 1, 2], dtype=dtype) + outputs = nn_ops.leaky_relu(constant_op.constant(np_values)) + with self.test_session() as sess: + outputs = sess.run(outputs) + tol = 2e-3 if dtype == np.float16 else 1e-6 + self.assertAllClose(outputs, [-0.4, -0.2, 0.0, 1.0, 2.0], rtol=tol, atol=tol) class SwishTest(test_lib.TestCase): diff --git a/tensorflow/python/ops/rnn_cell_impl.py b/tensorflow/python/ops/rnn_cell_impl.py index 1bf5551aff..f1ac3e9baf 100644 --- a/tensorflow/python/ops/rnn_cell_impl.py +++ b/tensorflow/python/ops/rnn_cell_impl.py @@ -988,6 +988,10 @@ class DropoutWrapper(RNNCell): return int(hashlib.md5(string).hexdigest()[:8], 16) & 0x7FFFFFFF @property + def wrapped_cell(self): + return self._cell + + @property def state_size(self): return self._cell.state_size diff --git a/tensorflow/python/summary/summary.py b/tensorflow/python/summary/summary.py index 355593eca5..92c1fcadd2 100644 --- a/tensorflow/python/summary/summary.py +++ b/tensorflow/python/summary/summary.py @@ -286,12 +286,13 @@ def merge(inputs, collections=None, name=None): return val -def merge_all(key=_ops.GraphKeys.SUMMARIES): +def merge_all(key=_ops.GraphKeys.SUMMARIES, scope=None): """Merges all summaries collected in the default graph. Args: key: `GraphKey` used to collect the summaries. Defaults to `GraphKeys.SUMMARIES`. + scope: Optional scope used to filter the summary ops, using `re.match` Returns: If no summaries were collected, returns None. Otherwise returns a scalar @@ -310,7 +311,7 @@ def merge_all(key=_ops.GraphKeys.SUMMARIES): raise RuntimeError( 'Merging tf.summary.* ops is not compatible with eager execution. ' 'Use tf.contrib.summary instead.') - summary_ops = _ops.get_collection(key) + summary_ops = _ops.get_collection(key, scope=scope) if not summary_ops: return None else: diff --git a/tensorflow/python/util/compat.py b/tensorflow/python/util/compat.py index 07382d93df..3ab0bd16fa 100644 --- a/tensorflow/python/util/compat.py +++ b/tensorflow/python/util/compat.py @@ -21,6 +21,7 @@ In addition to the functions below, `as_str` converts an object to a `str`. @@as_bytes @@as_text @@as_str_any +@@path_to_str ## Types The compatibility module also provides the following types: @@ -108,6 +109,20 @@ def as_str_any(value): return str(value) +def path_to_str(path): + """Returns the file system path representation of a `PathLike` object, else as it is. + + Args: + path: An object that can be converted to path representation. + + Returns: + A `str` object. + """ + if hasattr(path, "__fspath__"): + path = as_str_any(path.__fspath__()) + return path + + # Numpy 1.8 scalars don't inherit from numbers.Integral in Python 3, so we # need to check them specifically. The same goes from Real and Complex. integral_types = (_numbers.Integral, _np.integer) diff --git a/tensorflow/tools/api/golden/tensorflow.compat.pbtxt b/tensorflow/tools/api/golden/tensorflow.compat.pbtxt index ccc6031400..bab480ff9b 100644 --- a/tensorflow/tools/api/golden/tensorflow.compat.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.compat.pbtxt @@ -32,4 +32,8 @@ tf_module { name: "as_text" argspec: "args=[\'bytes_or_text\', \'encoding\'], varargs=None, keywords=None, defaults=[\'utf-8\'], " } + member_method { + name: "path_to_str" + argspec: "args=[\'path\'], varargs=None, keywords=None, defaults=None" + } } diff --git a/tensorflow/tools/api/golden/tensorflow.nn.rnn_cell.-dropout-wrapper.pbtxt b/tensorflow/tools/api/golden/tensorflow.nn.rnn_cell.-dropout-wrapper.pbtxt index f61a5a28e3..97edf245f6 100644 --- a/tensorflow/tools/api/golden/tensorflow.nn.rnn_cell.-dropout-wrapper.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.nn.rnn_cell.-dropout-wrapper.pbtxt @@ -88,6 +88,10 @@ tf_class { name: "weights" mtype: "<type \'property\'>" } + member { + name: "wrapped_cell" + mtype: "<type \'property\'>" + } member_method { name: "__init__" argspec: "args=[\'self\', \'cell\', \'input_keep_prob\', \'output_keep_prob\', \'state_keep_prob\', \'variational_recurrent\', \'input_size\', \'dtype\', \'seed\', \'dropout_state_filter_visitor\'], varargs=None, keywords=None, defaults=[\'1.0\', \'1.0\', \'1.0\', \'False\', \'None\', \'None\', \'None\', \'None\'], " diff --git a/tensorflow/tools/api/golden/tensorflow.pbtxt b/tensorflow/tools/api/golden/tensorflow.pbtxt index 35917e94ad..db1ed42185 100644 --- a/tensorflow/tools/api/golden/tensorflow.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.pbtxt @@ -1161,6 +1161,10 @@ tf_module { argspec: "args=[\'values\', \'value_range\', \'nbins\', \'dtype\', \'name\'], varargs=None, keywords=None, defaults=[\'100\', \"<dtype: \'int32\'>\", \'None\'], " } member_method { + name: "histogram_fixed_width_bins" + argspec: "args=[\'values\', \'value_range\', \'nbins\', \'dtype\', \'name\'], varargs=None, keywords=None, defaults=[\'100\', \"<dtype: \'int32\'>\", \'None\'], " + } + member_method { name: "identity" argspec: "args=[\'input\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], " } diff --git a/tensorflow/tools/api/golden/tensorflow.summary.pbtxt b/tensorflow/tools/api/golden/tensorflow.summary.pbtxt index 326e077d39..871ebb5247 100644 --- a/tensorflow/tools/api/golden/tensorflow.summary.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.summary.pbtxt @@ -50,7 +50,7 @@ tf_module { } member_method { name: "merge_all" - argspec: "args=[\'key\'], varargs=None, keywords=None, defaults=[\'summaries\'], " + argspec: "args=[\'key\', \'scope\'], varargs=None, keywords=None, defaults=[\'summaries\', \'None\'], " } member_method { name: "scalar" diff --git a/tensorflow/tools/benchmark/BUILD b/tensorflow/tools/benchmark/BUILD index caa6629c49..6ed2594e6a 100644 --- a/tensorflow/tools/benchmark/BUILD +++ b/tensorflow/tools/benchmark/BUILD @@ -61,10 +61,11 @@ tf_cc_test( # This binary may be built for either desktop or Android. # A typical Android build command will look like the following: -# bazel build -c opt tensorflow/core:android_tensorflow_lib \ +# bazel build tensorflow/core:android_tensorflow_lib \ # --crosstool_top=//external:android/crosstool \ # --cpu=armeabi-v7a \ # --host_crosstool_top=@bazel_tools//tools/cpp:toolchain +# --config monolithic tf_cc_binary( name = "benchmark_model", testonly = 1, diff --git a/tensorflow/tools/benchmark/README.md b/tensorflow/tools/benchmark/README.md index ca0da2d41b..e64af2bfe1 100644 --- a/tensorflow/tools/benchmark/README.md +++ b/tensorflow/tools/benchmark/README.md @@ -17,6 +17,7 @@ bazel build -c opt \ --crosstool_top=//external:android/crosstool \ --cpu=armeabi-v7a \ --host_crosstool_top=@bazel_tools//tools/cpp:toolchain \ + --config monolithic \ tensorflow/tools/benchmark:benchmark_model ``` diff --git a/tensorflow/tools/ci_build/builds/libtensorflow.sh b/tensorflow/tools/ci_build/builds/libtensorflow.sh index aadf480d37..9b3ff0cba7 100755 --- a/tensorflow/tools/ci_build/builds/libtensorflow.sh +++ b/tensorflow/tools/ci_build/builds/libtensorflow.sh @@ -51,7 +51,7 @@ function build_libtensorflow_tarball() { rm -rf ${DIR} TARBALL_SUFFIX="${1}" - BAZEL_OPTS="-c opt" + BAZEL_OPTS="-c opt --cxxopt=-D_GLIBCXX_USE_CXX11_ABI=0" export CC_OPT_FLAGS='-mavx' if [ "${TF_NEED_CUDA}" == "1" ]; then BAZEL_OPTS="${BAZEL_OPTS} --config=cuda" diff --git a/tensorflow/tools/ci_build/ci_sanity.sh b/tensorflow/tools/ci_build/ci_sanity.sh index b728c878da..aa341b144c 100755 --- a/tensorflow/tools/ci_build/ci_sanity.sh +++ b/tensorflow/tools/ci_build/ci_sanity.sh @@ -26,6 +26,8 @@ SCRIPT_DIR=$( cd ${0%/*} && pwd -P ) source "${SCRIPT_DIR}/builds/builds_common.sh" +ROOT_DIR=$( cd "$SCRIPT_DIR/../../.." && pwd -P ) + # Helper functions die() { echo $@ @@ -418,15 +420,8 @@ do_bazel_nobuild() { } do_pip_smoke_test() { - BUILD_CMD="bazel build ${BAZEL_FLAGS} //tensorflow/tools/pip_package:pip_smoke_test" - ${BUILD_CMD} - cmd_status \ - "Pip smoke test has failed. Please make sure any new TensorFlow are added to the tensorflow/tools/pip_package:build_pip_package dependencies." - - RUN_CMD="bazel-bin/tensorflow/tools/pip_package/pip_smoke_test" - ${RUN_CMD} - cmd_status \ - "The pip smoke test failed." + cd "$ROOT_DIR/tensorflow/tools/pip_package" + python pip_smoke_test.py } do_code_link_check() { @@ -500,20 +495,23 @@ do_clang_format_check() { } do_check_load_py_test() { - BUILD_CMD="bazel build ${BAZEL_FLAGS} //tensorflow/tools/pip_package:check_load_py_test" - ${BUILD_CMD} - cmd_status \ - "check_load_py_test failed to build." + cd "$ROOT_DIR/tensorflow/tools/pip_package" + python check_load_py_test.py +} - BUILD_CMD="bazel-bin/tensorflow/tools/pip_package/check_load_py_test" - ${BUILD_CMD} - cmd_status \ - "check_load_py_test failed." +do_cmake_python_sanity() { + cd "$ROOT_DIR/tensorflow/contrib/cmake" + python -m unittest -v python_sanity_test +} + +do_check_futures_test() { + cd "$ROOT_DIR/tensorflow/tools/test" + python check_futures_test.py } # Supply all sanity step commands and descriptions -SANITY_STEPS=("do_pylint PYTHON2" "do_pylint PYTHON3" "do_buildifier" "do_bazel_nobuild" "do_pip_package_licenses_check" "do_lib_package_licenses_check" "do_java_package_licenses_check" "do_pip_smoke_test" "do_check_load_py_test" "do_code_link_check") -SANITY_STEPS_DESC=("Python 2 pylint" "Python 3 pylint" "buildifier check" "bazel nobuild" "pip: license check for external dependencies" "C library: license check for external dependencies" "Java Native Library: license check for external dependencies" "Pip Smoke Test: Checking py_test dependencies exist in pip package" "Check load py_test: Check that BUILD files with py_test target properly load py_test" "Code Link Check: Check there are no broken links") +SANITY_STEPS=("do_pylint PYTHON2" "do_pylint PYTHON3" "do_check_futures_test" "do_buildifier" "do_bazel_nobuild" "do_pip_package_licenses_check" "do_lib_package_licenses_check" "do_java_package_licenses_check" "do_pip_smoke_test" "do_check_load_py_test" "do_code_link_check" "do_cmake_python_sanity") +SANITY_STEPS_DESC=("Python 2 pylint" "Python 3 pylint" "Check that python files have certain __future__ imports" "buildifier check" "bazel nobuild" "pip: license check for external dependencies" "C library: license check for external dependencies" "Java Native Library: license check for external dependencies" "Pip Smoke Test: Checking py_test dependencies exist in pip package" "Check load py_test: Check that BUILD files with py_test target properly load py_test" "Code Link Check: Check there are no broken links" "Test entries in /tensorflow/contrib/cmake/python_{modules|protos|protos_cc}.txt for validity and consistency") INCREMENTAL_FLAG="" DEFAULT_BAZEL_CONFIGS="--config=hdfs --config=gcp" @@ -548,7 +546,10 @@ while [[ ${COUNTER} -lt "${#SANITY_STEPS[@]}" ]]; do "${SANITY_STEPS[COUNTER]} (${SANITY_STEPS_DESC[COUNTER]}) ===" echo "" + # subshell: don't leak variables or changes of working directory + ( ${SANITY_STEPS[COUNTER]} ${INCREMENTAL_FLAG} + ) RESULT=$? if [[ ${RESULT} != "0" ]]; then diff --git a/tensorflow/tools/docker/parameterized_docker_build.sh b/tensorflow/tools/docker/parameterized_docker_build.sh index 1214b6b0a3..fa867b65db 100755 --- a/tensorflow/tools/docker/parameterized_docker_build.sh +++ b/tensorflow/tools/docker/parameterized_docker_build.sh @@ -414,7 +414,7 @@ if [[ ! -z "${TF_DOCKER_BUILD_PUSH_WITH_CREDENTIALS}" ]]; then if [[ $? != "0" ]]; then die "FAIL: Unable to login. Invalid credentials." fi - docker push $1 + docker push "${FINAL_IMG}" if [[ $? == "0" ]]; then docker logout echo "Successfully pushed Docker image ${FINAL_IMG}" diff --git a/tensorflow/tools/pip_package/BUILD b/tensorflow/tools/pip_package/BUILD index ff5dd6a0b0..c789e2ba0c 100644 --- a/tensorflow/tools/pip_package/BUILD +++ b/tensorflow/tools/pip_package/BUILD @@ -47,27 +47,6 @@ py_binary( deps = ["//tensorflow:tensorflow_py"], ) -py_test( - name = "pip_smoke_test", - srcs = ["pip_smoke_test.py"], - data = [ - "//tensorflow:all_opensource_files", - ], - tags = [ - "manual", - "notap", - ], -) - -py_binary( - name = "check_load_py_test", - srcs = ["check_load_py_test.py"], - data = [ - "//tensorflow:all_opensource_files", - ], - srcs_version = "PY2AND3", -) - # On Windows, python binary is a zip file of runfiles tree. # Add everything to its data dependency for generating a runfiles tree # for building the pip package on Windows. diff --git a/tensorflow/tools/pip_package/check_load_py_test.py b/tensorflow/tools/pip_package/check_load_py_test.py index 79d11b08ce..e2fe1121d7 100644 --- a/tensorflow/tools/pip_package/check_load_py_test.py +++ b/tensorflow/tools/pip_package/check_load_py_test.py @@ -22,6 +22,9 @@ import os import subprocess +os.chdir(os.path.abspath(os.path.join(os.path.dirname(__file__), '../../..'))) + + def check_output_despite_error(args): """Get output of args from command line, even if there are errors. diff --git a/tensorflow/tools/pip_package/pip_smoke_test.py b/tensorflow/tools/pip_package/pip_smoke_test.py index cddf9c8f44..8eee489e2d 100644 --- a/tensorflow/tools/pip_package/pip_smoke_test.py +++ b/tensorflow/tools/pip_package/pip_smoke_test.py @@ -23,8 +23,13 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function +import os import subprocess + +os.chdir(os.path.abspath(os.path.join(os.path.dirname(__file__), '../../..'))) + + PIP_PACKAGE_QUERY_EXPRESSION = \ 'deps(//tensorflow/tools/pip_package:build_pip_package)' @@ -134,8 +139,8 @@ def main(): raise RuntimeError("""One or more dependencies are not in the pip package. Please either blacklist the dependencies in -tensorflow/tensorflow/tensorflow/tools/pip_package/pip_smoke_test.py -or add them to tensorflow/tensorflow/tensorflow/tools/pip_package/BUILD.""") +//tensorflow/tools/pip_package/pip_smoke_test.py +or add them to //tensorflow/tools/pip_package/BUILD.""") else: print("TEST PASSED") diff --git a/tensorflow/tools/pip_package/setup.py b/tensorflow/tools/pip_package/setup.py index b2ca4a43c1..62df6453fb 100644 --- a/tensorflow/tools/pip_package/setup.py +++ b/tensorflow/tools/pip_package/setup.py @@ -29,7 +29,7 @@ from setuptools.dist import Distribution # This version string is semver compatible, but incompatible with pip. # For pip, we will remove all '-' characters from this string, and use the # result for pip. -_VERSION = '1.5.0-rc0' +_VERSION = '1.5.0-rc1' REQUIRED_PACKAGES = [ 'absl-py >= 0.1.6', diff --git a/tensorflow/tools/test/BUILD b/tensorflow/tools/test/BUILD index 28d651e910..159a8c1cfb 100644 --- a/tensorflow/tools/test/BUILD +++ b/tensorflow/tools/test/BUILD @@ -104,12 +104,3 @@ filegroup( ), visibility = ["//tensorflow:__subpackages__"], ) - -py_test( - name = "check_futures_test", - size = "small", - srcs = ["check_futures_test.py"], - data = ["//tensorflow:all_opensource_files"], - srcs_version = "PY2AND3", - deps = ["@six_archive//:six"], -) diff --git a/tensorflow/tools/test/check_futures_test.py b/tensorflow/tools/test/check_futures_test.py index 1c07511888..9181c9bd4a 100644 --- a/tensorflow/tools/test/check_futures_test.py +++ b/tensorflow/tools/test/check_futures_test.py @@ -33,7 +33,7 @@ import re import six -BASE_DIR = os.path.normpath(os.path.join(__file__, '../../..')) +BASE_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__), '../..')) FUTURES_PATTERN = re.compile(r'^from __future__ import (\w+)\s*$') FUTURES_PATTERN_2 = re.compile( r'^from __future__ import (\w+), (\w+), (\w+)\s*$') diff --git a/tensorflow/workspace.bzl b/tensorflow/workspace.bzl index d17dc81024..79e14e0e92 100644 --- a/tensorflow/workspace.bzl +++ b/tensorflow/workspace.bzl @@ -5,6 +5,7 @@ load("//third_party/mkl:build_defs.bzl", "mkl_repository") load("//third_party/git:git_configure.bzl", "git_configure") load("//third_party/py:python_configure.bzl", "python_configure") load("//third_party/sycl:sycl_configure.bzl", "sycl_configure") +load("//third_party/toolchains/clang6:repo.bzl", "clang6_configure") load("//third_party/toolchains/cpus/arm:arm_compiler_configure.bzl", "arm_compiler_configure") load("//third_party:repo.bzl", "tf_http_archive") load("@io_bazel_rules_closure//closure/private:java_import_external.bzl", "java_import_external") @@ -65,6 +66,7 @@ def tf_workspace(path_prefix="", tf_repo_name=""): # files, in case the parsing of those build files depends on the bazel # version we require here. check_bazel_version_at_least("0.5.4") + clang6_configure(name="local_config_clang6") cuda_configure(name="local_config_cuda") git_configure(name="local_config_git") sycl_configure(name="local_config_sycl") diff --git a/third_party/git/git_configure.bzl b/third_party/git/git_configure.bzl index 47e2125854..8e2839bdc2 100644 --- a/third_party/git/git_configure.bzl +++ b/third_party/git/git_configure.bzl @@ -38,6 +38,11 @@ def _git_conf_impl(repository_ctx): Label("@org_tensorflow//tensorflow/tools/git:gen_git_source.py")) generated_files_path = repository_ctx.path("gen") + r = repository_ctx.execute( + ["test", "-f", "%s/.git/logs/HEAD" % tensorflow_root_path]) + if r.return_code == 0: + unused_var = repository_ctx.path(Label("//:.git/HEAD")) # pylint: disable=unused-variable + result = repository_ctx.execute([ _get_python_bin(repository_ctx), python_script_path, "--configure", tensorflow_root_path, diff --git a/third_party/toolchains/clang6/BUILD b/third_party/toolchains/clang6/BUILD new file mode 100644 index 0000000000..ffd0fb0cdc --- /dev/null +++ b/third_party/toolchains/clang6/BUILD @@ -0,0 +1 @@ +package(default_visibility = ["//visibility:public"]) diff --git a/third_party/toolchains/clang6/CROSSTOOL.tpl b/third_party/toolchains/clang6/CROSSTOOL.tpl new file mode 100644 index 0000000000..6b7e5a8808 --- /dev/null +++ b/third_party/toolchains/clang6/CROSSTOOL.tpl @@ -0,0 +1,587 @@ +major_version: "v1" +minor_version: "llvm:6.0.0" +default_target_cpu: "k8" + +default_toolchain { + cpu: "k8" + toolchain_identifier: "k8-clang-6.0-cxx-4.8-linux-gnu" +} + +toolchain { + compiler: "clang6" # bazel build --compiler=clang6 + target_cpu: "k8" # bazel build --cpu=k8 + target_libc: "GLIBC_2.19" # bazel build --glibc=GLIBC_2.19 + + abi_libc_version: "2.19" + abi_version: "gcc-4.8-cxx11" + builtin_sysroot: "" + cc_target_os: "linux-gnu" + default_python_version: "python2.7" + dynamic_runtimes_filegroup: "dynamic-runtime-libs-k8" + host_system_name: "x86_64-unknown-linux-gnu" + needsPic: true + static_runtimes_filegroup: "static-runtime-libs-k8" + supports_embedded_runtimes: true + supports_fission: true + supports_gold_linker: true + supports_incremental_linker: true + supports_interface_shared_objects: true + supports_normalizing_ar: true + supports_start_end_lib: true + supports_thin_archives: true + target_system_name: "x86_64-unknown-linux-gnu" + toolchain_identifier: "k8-clang-6.0-cxx-4.8-linux-gnu" + + tool_path { name: "ar" path: "%package(@local_config_clang6//clang6)%/llvm/bin/llvm-ar" } + tool_path { name: "as" path: "%package(@local_config_clang6//clang6)%/llvm/bin/llvm-as" } + tool_path { name: "compat-ld" path: "%package(@local_config_clang6//clang6)%/llvm/bin/ld.lld" } + tool_path { name: "cpp" path: "%package(@local_config_clang6//clang6)%/llvm/bin/llvm-cpp" } + tool_path { name: "dwp" path: "%package(@local_config_clang6//clang6)%/llvm/bin/llvm-dwp" } + tool_path { name: "gcc" path: "%package(@local_config_clang6//clang6)%/llvm/bin/clang" } + tool_path { name: "gcov" path: "%package(@local_config_clang6//clang6)%/llvm/bin/llvm-cov" } + tool_path { name: "ld" path: "%package(@local_config_clang6//clang6)%/llvm/bin/ld.lld" } + tool_path { name: "llvm-profdata" path: "%package(@local_config_clang6//clang6)%/llvm/bin/llvm-profdata" } + tool_path { name: "nm" path: "%package(@local_config_clang6//clang6)%/llvm/bin/llvm-nm" } + tool_path { name: "objcopy" path: "%package(@local_config_clang6//clang6)%/llvm/bin/llvm-objcopy" } + tool_path { name: "objdump" path: "%package(@local_config_clang6//clang6)%/sbin/objdump" } + tool_path { name: "strip" path: "%package(@local_config_clang6//clang6)%/sbin/strip" } + + unfiltered_cxx_flag: "-no-canonical-prefixes" + + # Make C++ compilation deterministic. Use linkstamping instead of these + # compiler symbols. + unfiltered_cxx_flag: "-Wno-builtin-macro-redefined" + unfiltered_cxx_flag: "-D__DATE__=\"redacted\"" + unfiltered_cxx_flag: "-D__TIMESTAMP__=\"redacted\"" + unfiltered_cxx_flag: "-D__TIME__=\"redacted\"" + + objcopy_embed_flag: "-I" + objcopy_embed_flag: "binary" + + # This action_config makes features flags propagate + # to CC_FLAGS for genrules, and eventually skylark. + action_config { + action_name: "cc-flags-make-variable" + config_name: "cc-flags-make-variable" + } + + # Security hardening on by default. + # Conservative choice; -D_FORTIFY_SOURCE=2 may be unsafe in some cases. + # We need to undef it before redefining it as some distributions now have + # it enabled by default. + compiler_flag: "-U_FORTIFY_SOURCE" + compiler_flag: "-D_FORTIFY_SOURCE=1" + compiler_flag: "-fstack-protector" + linker_flag: "-Wl,-z,relro,-z,now" + + # This adds a little bit more durability to our Clang build. + # + # At the moment, this only only be needed for: + # - add_boringssl_s390x.patch: --Wa,--noexecstack + # + # Folks who do maintenance work on TF Bazel Clang should consider + # commenting out these lines, while doing that work, to gain a better + # understanding of what the intersection of support looks like between GCC + # and Clang. Please note that, Bazel does not support + # -Xclang-only / -Xgcc-only. + compiler_flag: "-Wno-unknown-warning-option" + compiler_flag: "-Wno-unused-command-line-argument" + compiler_flag: "-Wno-ignored-optimization-argument" + + #### Common compiler options. #### + compiler_flag: "-D_REENTRANT" + compiler_flag: "-D__STDC_FORMAT_MACROS" + compiler_flag: "-DSUPPRESS_USE_FILE_OFFSET64" + compiler_flag: "-Wall" + compiler_flag: "-Wformat-security" + compiler_flag: "-Wframe-larger-than=16384" + compiler_flag: "-Wno-char-subscripts" + compiler_flag: "-Wno-error=deprecated-declarations" + compiler_flag: "-Wno-uninitialized" + compiler_flag: "-Wno-sign-compare" + compiler_flag: "-Wno-strict-overflow" + compiler_flag: "-Wno-unused-function" + compiler_flag: "-fdiagnostics-show-option" + compiler_flag: "-fmessage-length=0" + compiler_flag: "-fno-exceptions" + compiler_flag: "-fno-omit-frame-pointer" + compiler_flag: "-fno-strict-aliasing" + compiler_flag: "-fno-use-init-array" + compiler_flag: "-funsigned-char" + compiler_flag: "-gmlt" + cxx_flag: "-Wno-deprecated" + cxx_flag: "-Wno-invalid-offsetof" # Needed for protobuf code (2017-11-07) + cxx_flag: "-fshow-overloads=best" + compiler_flag: "-Wthread-safety-analysis" + + # Python extensions unfortunately make this go wild. + compiler_flag: "-Wno-writable-strings" + + # GCC's warning produces too many false positives: + cxx_flag: "-Woverloaded-virtual" + cxx_flag: "-Wnon-virtual-dtor" + + # Enable coloring even if there's no attached terminal. Bazel removes the + # escape sequences if --nocolor is specified. This isn't supported by gcc + # on Ubuntu 14.04. + compiler_flag: "-fcolor-diagnostics" + + # Disable some broken warnings from Clang. + compiler_flag: "-Wno-ambiguous-member-template" + compiler_flag: "-Wno-pointer-sign" + + # These warnings have a low signal to noise ratio. + compiler_flag: "-Wno-reserved-user-defined-literal" + compiler_flag: "-Wno-return-type-c-linkage" + compiler_flag: "-Wno-invalid-source-encoding" + + # Per default we switch off any layering related warnings. + compiler_flag: "-Wno-private-header" + + # Clang-specific warnings that we explicitly enable for TensorFlow. Some of + # these aren't on by default, or under -Wall, or are subsets of warnings + # turned off above. + compiler_flag: "-Wfloat-overflow-conversion" + compiler_flag: "-Wfloat-zero-conversion" + compiler_flag: "-Wfor-loop-analysis" + compiler_flag: "-Wgnu-redeclared-enum" + compiler_flag: "-Winfinite-recursion" + compiler_flag: "-Wliteral-conversion" + compiler_flag: "-Wself-assign" + compiler_flag: "-Wstring-conversion" + compiler_flag: "-Wtautological-overlap-compare" + compiler_flag: "-Wunused-comparison" + compiler_flag: "-Wvla" + cxx_flag: "-Wdeprecated-increment-bool" + + # Clang code-generation flags for performance optimization. + compiler_flag: "-faligned-allocation" + compiler_flag: "-fnew-alignment=8" + + # Clang defaults to C99 while GCC defaults to C89. GCC plugins are written in + # C89 and don't have a BUILD rule we could add a copts flag to. + gcc_plugin_compiler_flag: "-std=gnu89" + + compilation_mode_flags { + mode: FASTBUILD + } + + compilation_mode_flags { + mode: DBG + compiler_flag: "-g" + } + + compilation_mode_flags { + mode: OPT + compiler_flag: "-g0" + compiler_flag: "-fdebug-types-section" + compiler_flag: "-DNDEBUG" + compiler_flag: "-fno-split-dwarf-inlining" + compiler_flag: "-Os" + compiler_flag: "-fexperimental-new-pass-manager" + compiler_flag: "-fdebug-info-for-profiling" + compiler_flag: "-ffunction-sections" + compiler_flag: "-fdata-sections" + linker_flag: "-Wl,--gc-sections" + linker_flag: "-Wl,-z,relro,-z,now" + } + + # Features indicating whether this is a host compile or not. Exactly one of + # these will be implicitly provided by bazel. + feature { name: "host" } + feature { name: "nonhost" } + + # Features indicating which compiler will be used for code generation. + feature { + name: "llvm_codegen" + provides: "codegen" + enabled: true + } + + # Features for compilation modes. Exactly one of these will be implicitly + # provided by bazel. + feature { name: "fastbuild" } + feature { name: "dbg" } + feature { name: "opt" } + + # Features controlling the C++ language mode. + feature { + name: "c++11" + provides: "c++std" + flag_set { + action: "c++-compile" + action: "c++-header-parsing" + action: "c++-header-preprocessing" + action: "c++-module-compile" + action: "linkstamp-compile" + flag_group { + flag: "-nostdinc++" + flag: "-std=c++11" + flag: "-Wc++14-extensions" + flag: "-Wc++2a-extensions" + flag: "-Wno-binary-literal" + } + } + } + feature { + name: "c++14" + provides: "c++std" + flag_set { + action: "c++-compile" + action: "c++-header-parsing" + action: "c++-header-preprocessing" + action: "c++-module-compile" + action: "linkstamp-compile" + flag_group { + flag: "-nostdinc++" + flag: "-std=c++14" + flag: "-Wc++11-compat" + flag: "-Wno-c++11-compat-binary-literal" + flag: "-Wc++2a-extensions" + } + } + } + feature { + name: "c++17" + provides: "c++std" + flag_set { + action: "c++-compile" + action: "c++-header-parsing" + action: "c++-header-preprocessing" + action: "c++-module-compile" + action: "linkstamp-compile" + flag_group { + flag: "-nostdinc++" + flag: "-std=c++17" + flag: "-Wc++11-compat" + flag: "-Wno-c++11-compat-binary-literal" + flag: "-Wc++2a-extensions" + } + } + } + feature { + name: "c++2a" + provides: "c++std" + flag_set { + action: "c++-compile" + action: "c++-header-parsing" + action: "c++-header-preprocessing" + action: "c++-module-compile" + action: "linkstamp-compile" + flag_group { + flag: "-nostdinc++" + flag: "-std=c++2a" + flag: "-Wc++11-compat" + flag: "-Wno-c++11-compat-binary-literal" + } + } + } + feature { + name: "c++default" + enabled: true + flag_set { + # Provide the c++11 flags if no standard is selected + with_feature { + not_feature: "c++11" + not_feature: "c++14" + not_feature: "c++17" + not_feature: "c++2a" + } + action: "c++-compile" + action: "c++-header-parsing" + action: "c++-header-preprocessing" + action: "c++-module-compile" + action: "linkstamp-compile" + flag_group { + flag: "-nostdinc++" + flag: "-std=c++11" + flag: "-Wc++14-extensions" + flag: "-Wc++2a-extensions" + flag: "-Wno-binary-literal" + } + } + } + + feature { + name: "use_compiler_rt" + requires { feature: "llvm_codegen" } + # TODO(saugustine): At the moment, "use_compiler_rt" also + # requires "linking_mode_flags { mode: FULLY_STATIC" ... }, + # but that isn't a feature. We should probably convert it. + flag_set { + action: "c++-link" + action: "c++-link-interface-dynamic-library" + action: "c++-link-dynamic-library" + action: "c++-link-executable" + # "link" is a misnomer for these actions. They are really just + # invocations of ar. + #action: "c++-link-pic-static-library" + #action: "c++-link-static-library" + #action: "c++-link-alwayslink-static-library" + #action: "c++-link-pic-static-library" + #action: "c++-link-alwayslink-pic-static-library" + flag_group { + flag: "-rtlib=compiler-rt" + flag: "-lunwind" + } + } + } + + feature { + name: "pie" + flag_set { + action: "assemble" + action: "preprocess-assemble" + action: "c-compile" + action: "c++-compile" + action: "c++-header-parsing" + action: "c++-header-preprocessing" + action: "c++-module-compile" + action: "c++-module-codegen" + action: "cc-flags-make-variable" + action: "lto-backend" + action: "linkstamp-compile" + flag_group { + flag: "-mpie-copy-relocations" + flag: "-fPIE" + } + } + flag_set { + action: "cc-flags-make-variable" + action: "c++-link-executable" + flag_group { + flag: "-pie" + } + } + } + + # Pic must appear after pie, because pic may need to override pie, and bazel + # turns it on selectively. These don't interact with other options. + # + # TODO: In practice, normal vs pic vs pie is a ternary mode. We should + # implement it that way. This will require changes to bazel, which only + # calculates whether or not pic is needed, not pie. + # + # NOTE: Bazel might make this all a moot point. + feature { + name: "pic" + flag_set { + action: "assemble" + action: "preprocess-assemble" + action: "c-compile" + action: "c++-compile" + action: "c++-module-codegen" + action: "c++-module-compile" + action: "linkstamp-compile" + expand_if_all_available: "pic" + flag_group { + flag: "-fPIC" + } + } + } + + feature { + name: "gold" + enabled: true + flag_set { + action: "c++-link-executable" + action: "c++-link-dynamic-library" + action: "c++-link-interface-dynamic-library" + flag_group { + expand_if_none_available: "lto" + flag: "-fuse-ld=gold" + } + } + } + + # This is great if you want linking TensorFlow to take ten minutes. + feature { + name: "lto" + requires { feature: "nonhost" } + flag_set { + action: "c-compile" + action: "c++-compile" + flag_group { + flag: "-flto=thin" + } + } + flag_set { + action: "c++-link-executable" + action: "c++-link-dynamic-library" + action: "c++-link-interface-dynamic-library" + flag_group { + flag: "-flto=thin" + } + } + } + + feature { + name: "parse_headers" + flag_set { + action: "c++-header-parsing" + flag_group { + flag: "-xc++-header" + flag: "-fsyntax-only" + } + } + } + + feature { + name: "preprocess_headers" + flag_set { + action: "c++-header-preprocessing" + flag_group { + flag: "-xc++" + flag: "-E" + } + } + } + + feature { + name: "per_object_debug_info" + flag_set { + action: "c-compile" + action: "c++-compile" + action: "c++-module-codegen" + action: "assemble" + action: "preprocess-assemble" + action: "lto-backend" + flag_group { + flag: "-gsplit-dwarf" + flag: "-ggnu-pubnames" + } + } + flag_set { + action: "c++-link-executable" + action: "c++-link-dynamic-library" + action: "c++-link-interface-dynamic-library" + flag_group { + expand_if_all_available: "is_using_fission" + flag: "-Wl,--gdb-index" + } + } + } + + feature { + name: "xray" + requires { + feature: "llvm_codegen" + feature: "nonhost" + } + flag_set { + action: "c-compile" + action: "c++-compile" + action: "c++-header-parsing" + action: "c++-header-preprocessing" + action: "c++-module-compile" + action: "c++-link-interface-dynamic-library" + action: "c++-link-dynamic-library" + action: "c++-link-executable" + flag_group { + flag: "-fxray-instrument" + } + } + } + + feature { + name: "minimal_ubsan" + requires { feature: "llvm_codegen" } + flag_set { + action: "c-compile" + action: "c++-compile" + action: "c++-header-parsing" + action: "c++-header-preprocessing" + action: "c++-module-compile" + action: "c++-module-codegen" + flag_group { + flag: "-fsanitize=return,returns-nonnull-attribute,vla-bound,unreachable,float-cast-overflow" + flag: "-fsanitize-trap=all" + flag: "-DUNDEFINED_BEHAVIOR_SANITIZER" + } + } + } + + feature { + name: "minimal_ubsan_enabled_by_default" + requires { + feature: "llvm_codegen" + feature: "fastbuild" + } + enabled: true + implies: "minimal_ubsan" + } + + cxx_builtin_include_directory: "%package(@local_config_clang6//clang6)%/llvm/lib/clang/6.0.0/include" + cxx_builtin_include_directory: "/usr/include" + + unfiltered_cxx_flag: "-cxx-isystem" + unfiltered_cxx_flag: "/usr/include/c++/4.8" + unfiltered_cxx_flag: "-cxx-isystem" + unfiltered_cxx_flag: "/usr/include/x86_64-linux-gnu/c++/4.8" + unfiltered_cxx_flag: "-isystem" + unfiltered_cxx_flag: "%package(@local_config_clang6//clang6)%/llvm/lib/clang/6.0.0/include" + unfiltered_cxx_flag: "-isystem" + unfiltered_cxx_flag: "/usr/include/x86_64-linux-gnu" + unfiltered_cxx_flag: "-isystem" + unfiltered_cxx_flag: "/usr/include" + + linker_flag: "-Wl,--build-id=md5" + linker_flag: "-Wl,--fatal-warnings" + linker_flag: "-Wl,--hash-style=gnu" + linker_flag: "-no-canonical-prefixes" + linker_flag: "--target=x86_64-unknown-linux-gnu" + + linker_flag: "-L/usr/lib/gcc/x86_64-linux-gnu/4.8" + + # This is the minimum x86 architecture TensorFlow supports. + compiler_flag: "-DARCH_K8" + compiler_flag: "-m64" + + # These are for Linux. + ld_embed_flag: "-melf_x86_64" + linker_flag: "-Wl,--eh-frame-hdr" + linker_flag: "-Wl,-z,max-page-size=0x1000" + + # Google never uses the stack like a heap, e.g. alloca(), because tcmalloc + # and jemalloc are so fast. However copts=["$(STACK_FRAME_UNLIMITED)"] can be + # specified when that can't be the case. + make_variable { + name: "STACK_FRAME_UNLIMITED" + value: "-Wframe-larger-than=100000000 -Wno-vla" + } + + # These flags are for folks who build C/C++ code inside genrules. + make_variable { + name: "CC_FLAGS" + value: "-no-canonical-prefixes --target=x86_64-unknown-linux-gnu -fno-omit-frame-pointer -fno-tree-vrp -msse3" + } + + feature { + name: "copts" + flag_set { + expand_if_all_available: "copts" + action: "assemble" + action: "preprocess-assemble" + action: "c-compile" + action: "c++-compile" + action: "c++-header-parsing" + action: "c++-header-preprocessing" + action: "c++-module-compile" + action: "c++-module-codegen" + action: "lto-backend" + flag_group { + iterate_over: "copts" + flag: "%{copts}" + } + } + } + + # Please do not statically link libstdc++. This would probably lead to a lot + # of bloat since OpKernels need to use linkstatic=1 because b/27630669 and + # it could cause memory leaks since Python uses dlopen() on our libraries: + # https://stackoverflow.com/a/35015415 + linker_flag: "-lstdc++" + linker_flag: "-lm" + linker_flag: "-lpthread" + linker_flag: "-l:/lib/x86_64-linux-gnu/libc-2.19.so" +} diff --git a/third_party/toolchains/clang6/README.md b/third_party/toolchains/clang6/README.md new file mode 100644 index 0000000000..0c6be25a0e --- /dev/null +++ b/third_party/toolchains/clang6/README.md @@ -0,0 +1,101 @@ +# TensorFlow Bazel Clang + +This is a specialized toolchain that uses an old Debian with a new Clang that +can cross compile to any x86_64 microarchitecture. It's intended to build Linux +binaries that only require the following ABIs: + +- GLIBC_2.18 +- CXXABI_1.3.7 (GCC 4.8.3) +- GCC_4.2.0 + +Which are available on at least the following Linux platforms: + +- Ubuntu 14+ +- CentOS 7+ +- Debian 8+ +- SuSE 13.2+ +- Mint 17.3+ +- Manjaro 0.8.11 + +# System Install + +On Debian 8 (Jessie) Clang 6.0 can be installed as follows: + +```sh +cat >>/etc/apt/sources.list <<'EOF' +deb http://apt.llvm.org/jessie/ llvm-toolchain-jessie main +deb-src http://apt.llvm.org/jessie/ llvm-toolchain-jessie main +EOF +wget -O - https://apt.llvm.org/llvm-snapshot.gpg.key | apt-key add - +apt-key fingerprint |& grep '6084 F3CF 814B 57C1 CF12 EFD5 15CF 4D18 AF4F 7421' +apt-get update +apt-get install clang lld +``` + +# Bazel Configuration + +This toolchain can compile TensorFlow in 2m30s on a 96-core Skylake GCE VM if +the following `.bazelrc` settings are added: + +``` +startup --host_jvm_args=-Xmx30G +startup --host_jvm_args=-Xms30G +startup --host_jvm_args=-XX:MaxNewSize=3g +startup --host_jvm_args=-XX:-UseAdaptiveSizePolicy +startup --host_jvm_args=-XX:+UseConcMarkSweepGC +startup --host_jvm_args=-XX:TargetSurvivorRatio=70 +startup --host_jvm_args=-XX:SurvivorRatio=6 +startup --host_jvm_args=-XX:+UseCMSInitiatingOccupancyOnly +startup --host_jvm_args=-XX:CMSFullGCsBeforeCompaction=1 +startup --host_jvm_args=-XX:CMSInitiatingOccupancyFraction=75 + +build --jobs=100 +build --local_resources=200000,100,100 +build --crosstool_top=@local_config_clang6//clang6 +build --noexperimental_check_output_files +build --nostamp +build --config=opt +build --noexperimental_check_output_files +build --copt=-march=native +build --host_copt=-march=native +``` + +# x86_64 Microarchitectures + +## Intel CPU Line + +- 2003 P6 M SSE SSE2 +- 2004 prescott SSE3 SSSE3 (-march=prescott) +- 2006 core X64 SSE4.1 (only on 45nm variety) (-march=core2) +- 2008 nehalem SSE4.2 VT-x VT-d (-march=nehalem) +- 2010 westmere CLMUL AES (-march=westmere) +- 2012 sandybridge AVX TXT (-march=sandybridge) +- 2012 ivybridge F16C MOVBE (-march=ivybridge) +- 2013 haswell AVX2 TSX BMI2 FMA (-march=haswell) +- 2014 broadwell RDSEED ADCX PREFETCHW (-march=broadwell - works on trusty gcc4.9) +- 2015 skylake SGX ADX MPX AVX-512[xeon-only] (-march=skylake / -march=skylake-avx512 - needs gcc7) +- 2018 cannonlake AVX-512 SHA (-march=cannonlake - needs clang5) + +## Intel Low Power CPU Line + +- 2013 silvermont SSE4.1 SSE4.2 VT-x (-march=silvermont) +- 2016 goldmont SHA (-march=goldmont - needs clang5) + +## AMD CPU Line + +- 2003 k8 SSE SSE2 (-march=k8) +- 2005 k8 (Venus) SSE3 (-march=k8-sse3) +- 2008 barcelona SSE4a?! (-march=barcelona) +- 2011 bulldozer SSE4.1 SSE4.2 CLMUL AVX AES FMA4?! (-march=bdver1) +- 2011 piledriver FMA (-march=bdver2) +- 2015 excavator AVX2 BMI2 MOVBE (-march=bdver4) + +## Google Compute Engine Supported CPUs + +- 2012 sandybridge 2.6gHz -march=sandybridge +- 2012 ivybridge 2.5gHz -march=ivybridge +- 2013 haswell 2.3gHz -march=haswell +- 2014 broadwell 2.2gHz -march=broadwell +- 2015 skylake 2.0gHz -march=skylake-avx512 + +See: <https://cloud.google.com/compute/docs/cpu-platforms> diff --git a/third_party/toolchains/clang6/clang.BUILD b/third_party/toolchains/clang6/clang.BUILD new file mode 100644 index 0000000000..802d62c17c --- /dev/null +++ b/third_party/toolchains/clang6/clang.BUILD @@ -0,0 +1,162 @@ +package(default_visibility = ["//visibility:public"]) + +# Please note that the output of these tools is unencumbered. +licenses(["restricted"]) # NCSA, GPLv3 (e.g. gold) + +filegroup( + name = "ar", + srcs = ["llvm/bin/llvm-ar"], + output_licenses = ["unencumbered"], +) + +filegroup( + name = "as", + srcs = ["llvm/bin/llvm-as"], + output_licenses = ["unencumbered"], +) + +filegroup( + name = "cpp", + srcs = ["llvm/bin/llvm-cpp"], + output_licenses = ["unencumbered"], +) + +filegroup( + name = "dwp", + srcs = ["llvm/bin/llvm-dwp"], + output_licenses = ["unencumbered"], +) + +filegroup( + name = "gcc", + srcs = ["llvm/bin/clang"], + output_licenses = ["unencumbered"], +) + +filegroup( + name = "gcov", + srcs = ["llvm/bin/llvm-cov"], + output_licenses = ["unencumbered"], +) + +filegroup( + name = "ld", + srcs = ["llvm/bin/ld.lld"], + output_licenses = ["unencumbered"], +) + +filegroup( + name = "nm", + srcs = ["llvm/bin/llvm-nm"], + output_licenses = ["unencumbered"], +) + +filegroup( + name = "objcopy", + srcs = ["llvm/bin/llvm-objcopy"], + output_licenses = ["unencumbered"], +) + +filegroup( + name = "objdump", + srcs = ["llvm/bin/llvm-objdump"], + output_licenses = ["unencumbered"], +) + +filegroup( + name = "profdata", + srcs = ["llvm/bin/llvm-profdata"], + output_licenses = ["unencumbered"], +) + +filegroup( + name = "strip", + srcs = ["sbin/strip"], + output_licenses = ["unencumbered"], +) + +filegroup( + name = "xray", + srcs = ["llvm/bin/llvm-xray"], + output_licenses = ["unencumbered"], +) + +filegroup( + name = "includes", + srcs = glob(["llvm/lib/clang/6.0.0/include/**"]), + output_licenses = ["unencumbered"], +) + +filegroup( + name = "libraries", + srcs = glob([ + "lib/*.*", + "lib/clang/6.0.0/lib/linux/*.*", + ]), + output_licenses = ["unencumbered"], +) + +filegroup( + name = "compiler_files", + srcs = [ + ":as", + ":gcc", + ":includes", + ], + output_licenses = ["unencumbered"], +) + +filegroup( + name = "linker_files", + srcs = [ + ":ar", + ":ld", + ":libraries", + ], + output_licenses = ["unencumbered"], +) + +filegroup( + name = "all_files", + srcs = [ + ":compiler_files", + ":dwp", + ":gcov", + ":linker_files", + ":nm", + ":objcopy", + ":objdump", + ":profdata", + ":strip", + ":xray", + ], + output_licenses = ["unencumbered"], +) + +filegroup( + name = "empty", + srcs = [], # bazel crashes without this + output_licenses = ["unencumbered"], +) + +cc_toolchain_suite( + name = "clang6", + toolchains = { + "k8|clang6": ":clang6-k8", + }, +) + +cc_toolchain( + name = "clang6-k8", + all_files = ":all_files", + compiler_files = ":compiler_files", + cpu = "k8", + dwp_files = ":dwp", + dynamic_runtime_libs = [":empty"], + linker_files = ":linker_files", + objcopy_files = ":objcopy", + output_licenses = ["unencumbered"], + static_runtime_libs = [":empty"], + strip_files = ":strip", + supports_param_files = 1, +) diff --git a/third_party/toolchains/clang6/repo.bzl b/third_party/toolchains/clang6/repo.bzl new file mode 100644 index 0000000000..b81f44506f --- /dev/null +++ b/third_party/toolchains/clang6/repo.bzl @@ -0,0 +1,30 @@ +"""Repository rule for Debian 8 Jessie Clang-6.0 portable Linux builds.""" + +def _clang6_configure(ctx): + # TODO(jart): It'd probably be better to use Bazel's struct.to_proto() + # method to generate a gigantic CROSSTOOL file that allows + # Clang to support everything. + ctx.symlink( + ctx.os.environ.get('TF_LLVM_PATH', + '/usr/lib/llvm-6.0'), + 'clang6/llvm') + ctx.symlink( + ctx.os.environ.get('STRIP', '/usr/bin/strip'), + 'clang6/sbin/strip') + ctx.symlink( + ctx.os.environ.get('OBJDUMP', '/usr/bin/objdump'), + 'clang6/sbin/objdump') + ctx.symlink(ctx.attr._build, 'clang6/BUILD') + ctx.template('clang6/CROSSTOOL', ctx.attr._crosstool, { + '%package(@local_config_clang6//clang6)%': str(ctx.path('clang6')), + }) + +clang6_configure = repository_rule( + implementation = _clang6_configure, + attrs = { + '_build': attr.label( + default=str(Label('//third_party/toolchains/clang6:clang.BUILD'))), + '_crosstool': attr.label( + default=str(Label('//third_party/toolchains/clang6:CROSSTOOL.tpl'))), + }, +) |
