[SE] Set denormal and rounding modes in HostStream thread.

Previously, HostStream used a TF ThreadPool, which sets denormal and rounding modes. A previous change switched from using a ThreadPool to using a Thread, which does not set denormal and rounding modes. This change restores the previous denormal/rounding behavior to unbreak tensorflow/compiler/xla/tests:exhaustive_op_test_cpu, pending a more complete fix.

PiperOrigin-RevId: 247064380
This commit is contained in:
Peter Hawkins 2019-05-07 12:02:50 -07:00 committed by TensorFlower Gardener
parent 581577abf5
commit 86b2e49082
2 changed files with 8 additions and 0 deletions

View File

@ -67,6 +67,7 @@ cc_library(
"host_stream.h",
],
deps = [
"//tensorflow/core:lib_internal",
"//tensorflow/stream_executor:kernel",
"//tensorflow/stream_executor/lib",
"@com_google_absl//absl/synchronization",

View File

@ -18,6 +18,8 @@ limitations under the License.
#include "tensorflow/stream_executor/host/host_stream.h"
#include "absl/synchronization/notification.h"
#include "tensorflow/core/platform/denormal.h"
#include "tensorflow/core/platform/setround.h"
namespace stream_executor {
namespace host {
@ -45,6 +47,11 @@ bool HostStream::EnqueueTask(std::function<void()> fn) {
bool HostStream::WorkAvailable() { return !work_queue_.empty(); }
void HostStream::WorkLoop() {
// Set denormal and rounding behavior to match the default TF ThreadPool
// behavior.
// TODO(phawkins, jlebar): it's not clear this is the best place to set this.
tensorflow::port::ScopedFlushDenormal flush;
tensorflow::port::ScopedSetRound round(FE_TONEAREST);
while (true) {
std::function<void()> fn;
{