[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:
parent
581577abf5
commit
86b2e49082
@ -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",
|
||||
|
@ -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;
|
||||
{
|
||||
|
Loading…
Reference in New Issue
Block a user