diff --git a/tensorflow/stream_executor/host/BUILD b/tensorflow/stream_executor/host/BUILD index 00fabe5772f..6ad06bb9bb9 100644 --- a/tensorflow/stream_executor/host/BUILD +++ b/tensorflow/stream_executor/host/BUILD @@ -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", diff --git a/tensorflow/stream_executor/host/host_stream.cc b/tensorflow/stream_executor/host/host_stream.cc index 0d8cb46f196..413edc6739a 100644 --- a/tensorflow/stream_executor/host/host_stream.cc +++ b/tensorflow/stream_executor/host/host_stream.cc @@ -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 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 fn; {