add denormal options (ftz and daz)

Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/15423

Reviewed By: yinghai

Differential Revision: D13526340

fbshipit-source-id: de2ecc717b4f778f33a8bf940ed144dbb230c7a8

diff --git a/caffe2/core/init_denormals.cc b/caffe2/core/init_denormals.cc
new file mode 100644 (file)
index 0000000..f938b47
--- /dev/null
+++ b/caffe2/core/init_denormals.cc
@@ -0,0 +1,38 @@
+#if defined(__SSE3__)
+
+#include <immintrin.h>
+
+#include "caffe2/core/common.h"
+#include "caffe2/core/init.h"
+
+C10_DEFINE_int(
+    caffe2_ftz,
+    false,
+    "If true, turn on flushing denormals to zero (FTZ)");
+C10_DEFINE_int(
+    caffe2_daz,
+    false,
+    "If true, turn on replacing denormals loaded from memory with zero (DAZ)");
+
+namespace caffe2 {
+
+bool Caffe2SetDenormals(int*, char***) {
+  if (FLAGS_caffe2_ftz) {
+    VLOG(1) << "Setting FTZ";
+    _MM_SET_FLUSH_ZERO_MODE(_MM_FLUSH_ZERO_ON);
+  }
+  if (FLAGS_caffe2_daz) {
+    VLOG(1) << "Setting DAZ";
+    _MM_SET_DENORMALS_ZERO_MODE(_MM_DENORMALS_ZERO_ON);
+  }
+  return true;
+}
+
+REGISTER_CAFFE2_INIT_FUNCTION(
+    Caffe2SetDenormals,
+    &Caffe2SetDenormals,
+    "Set denormal settings.");
+
+} // namespace caffe2
+
+#endif