Disables denormal floating numbers on ARM CPU

aten/src/ATen/cpu/FlushDenormal.cpp

@@ -28,6 +28,45 @@ bool set_flush_denormal(bool on) {
   }
   return false;
 }
+#elif defined(__ARM_FP) && (__ARM_FP > 0)
+// Imported from TensorFlow, tensorflow/third_party/xla/third_party/tsl/tsl/platform/denormal.cc
+// Copyright 2015 The TensorFlow Authors. All Rights Reserved.
+
+// Flush-to-zero bit on the ARM floating-point control register.
+#define ARM_FPCR_FZ   (1 << 24)
+
+static inline void ArmSetFloatingPointControlRegister(uint32_t fpcr) {
+#if defined(__aarch64__)
+  __asm__ __volatile__("msr fpcr, %[fpcr]"
+                       :
+                       : [fpcr] "r"(static_cast<uint64_t>(fpcr)));
+#else
+  __asm__ __volatile__("vmsr fpscr, %[fpcr]" : : [fpcr] "r"(fpcr));
+#endif
+}
+
+static inline uint32_t ArmGetFloatingPointControlRegister() {
+  uint32_t fpcr;
+#if defined(__aarch64__)
+  uint64_t fpcr64;
+  __asm__ __volatile__("mrs %[fpcr], fpcr" : [fpcr] "=r"(fpcr64));
+  fpcr = static_cast<uint32_t>(fpcr64);
+#else
+  __asm__ __volatile__("vmrs %[fpcr], fpscr" : [fpcr] "=r"(fpcr));
+#endif
+  return fpcr;
+}
+
+bool set_flush_denormal(bool on) {
+    uint32_t fpcr = ArmGetFloatingPointControlRegister();
+    if (on) {
+      fpcr |= ARM_FPCR_FZ;
+    } else {
+      fpcr &= ~ ARM_FPCR_FZ;
+    }
+    ArmSetFloatingPointControlRegister(fpcr);
+    return true;
+}
 #else
 bool set_flush_denormal(bool on) {
   return false;