add scoped flush-to-zero and denormals-are-zero flags

lib/util_cpu.h

@@ -21,6 +21,22 @@
 
 namespace qsim {
 
+// Gets the current MXCSR control flags.
+inline unsigned GetMXCSRControlFlags() {
+#ifdef __SSE2__
+  return _mm_getcsr();
+#else
+  return 0;
+#endif
+}
+
+// Sets the MXCSR control flags to a specific value.
+inline void SetMXCSRControlFlags(unsigned value) {
+#ifdef __SSE2__
+  _mm_setcsr(value);
+#endif
+}
+
 // This function sets flush-to-zero and denormals-are-zeros MXCSR control
 // flags. This prevents rare cases of performance slowdown potentially at
 // the cost of a tiny precision loss.
@@ -38,6 +54,38 @@ inline void ClearFlushToZeroAndDenormalsAreZeros() {
 #endif
 }
 
+// RAII guard for the flush-to-zero and denormals-are-zeros MXCSR control flags.
+class ScopedFlushToZeroAndDenormalsAreZeros {
+ public:
+  explicit ScopedFlushToZeroAndDenormalsAreZeros(bool denormals_are_zeros)
+      : original_flags_(GetMXCSRControlFlags()) {
+    if (denormals_are_zeros) {
+      SetFlushToZeroAndDenormalsAreZeros();
+    } else {
+      ClearFlushToZeroAndDenormalsAreZeros();
+    }
+  }
+
+  ~ScopedFlushToZeroAndDenormalsAreZeros() {
+    SetMXCSRControlFlags(original_flags_);
+  }
+
+  // Prevent copying.
+  ScopedFlushToZeroAndDenormalsAreZeros(
+      const ScopedFlushToZeroAndDenormalsAreZeros&) = delete;
+  ScopedFlushToZeroAndDenormalsAreZeros& operator=(
+      const ScopedFlushToZeroAndDenormalsAreZeros&) = delete;
+
+  // Prevent moving
+  ScopedFlushToZeroAndDenormalsAreZeros(
+      ScopedFlushToZeroAndDenormalsAreZeros&&) = delete;
+  ScopedFlushToZeroAndDenormalsAreZeros& operator=(
+      ScopedFlushToZeroAndDenormalsAreZeros&&) = delete;
+
+ private:
+  unsigned original_flags_;
+};
+
 }  // namespace qsim
 
 #endif  // UTIL_CPU_H_

pybind_interface/cuda/pybind_main_cuda.cpp

@@ -50,8 +50,15 @@ namespace qsim {
     StateSpace::Parameter ss_params;
   };
 
+  inline void GetMXCSRControlFlags() {}
+  inline void SetMXCSRControlFlags(unsigned value) {}
   inline void SetFlushToZeroAndDenormalsAreZeros() {}
   inline void ClearFlushToZeroAndDenormalsAreZeros() {}
+
+  class ScopedFlushToZeroAndDenormalsAreZeros {
+   public:
+    explicit ScopedFlushToZeroAndDenormalsAreZeros(bool denormals_are_zeros) {}
+  };
 }
 
 #include "../pybind_main.cpp"

pybind_interface/custatevec/pybind_main_custatevec.cpp

@@ -59,8 +59,15 @@ namespace qsim {
     custatevecHandle_t custatevec_handle;
   };
 
+  inline void GetMXCSRControlFlags() {}
+  inline void SetMXCSRControlFlags(unsigned value) {}
   inline void SetFlushToZeroAndDenormalsAreZeros() {}
   inline void ClearFlushToZeroAndDenormalsAreZeros() {}
+
+  class ScopedFlushToZeroAndDenormalsAreZeros {
+   public:
+    explicit ScopedFlushToZeroAndDenormalsAreZeros(bool denormals_are_zeros) {}
+  };
 }
 
 #include "../pybind_main.cpp"

pybind_interface/custatevecex/pybind_main_custatevecex.cpp

@@ -103,8 +103,15 @@ namespace qsim {
     unsigned nwt = 0;
   };
 
+  inline void GetMXCSRControlFlags() {}
+  inline void SetMXCSRControlFlags(unsigned value) {}
   inline void SetFlushToZeroAndDenormalsAreZeros() {}
   inline void ClearFlushToZeroAndDenormalsAreZeros() {}
+
+  class ScopedFlushToZeroAndDenormalsAreZeros {
+   public:
+    explicit ScopedFlushToZeroAndDenormalsAreZeros(bool denormals_are_zeros) {}
+  };
 }
 
 #include "../pybind_main.cpp"

pybind_interface/hip/pybind_main_hip.cpp

@@ -50,8 +50,15 @@ namespace qsim {
     StateSpace::Parameter ss_params;
   };
 
+  inline void GetMXCSRControlFlags() {}
+  inline void SetMXCSRControlFlags(unsigned value) {}
   inline void SetFlushToZeroAndDenormalsAreZeros() {}
   inline void ClearFlushToZeroAndDenormalsAreZeros() {}
+
+  class ScopedFlushToZeroAndDenormalsAreZeros {
+   public:
+    explicit ScopedFlushToZeroAndDenormalsAreZeros(bool denormals_are_zeros) {}
+  };
 }
 
 #include "../pybind_main.cpp"

pybind_interface/pybind_main.cpp

@@ -349,11 +349,7 @@ std::vector<std::complex<float>> qsim_simulate(const py::dict &options) {
     param.verbosity = ParseOptions<unsigned>(options, "v\0");
     param.seed = ParseOptions<unsigned>(options, "s\0");
 
-    if (denormals_are_zeros) {
-      SetFlushToZeroAndDenormalsAreZeros();
-    } else {
-      ClearFlushToZeroAndDenormalsAreZeros();
-    }
+    ScopedFlushToZeroAndDenormalsAreZeros mxcsr_guard(denormals_are_zeros);
 
     Factory factory(options);
     Runner::Run(param, factory, circuit, measure);
@@ -409,11 +405,7 @@ std::vector<std::complex<float>> qtrajectory_simulate(const py::dict &options) {
       }
     };
 
-    if (denormals_are_zeros) {
-      SetFlushToZeroAndDenormalsAreZeros();
-    } else {
-      ClearFlushToZeroAndDenormalsAreZeros();
-    }
+    ScopedFlushToZeroAndDenormalsAreZeros mxcsr_guard(denormals_are_zeros);
 
     if (!NoisyRunner::RunBatch(param, ncircuit, seed, seed + 1, state_space,
                                simulator, measure)) {
@@ -595,7 +587,6 @@ class SimulatorHelper {
       : factory(Factory(options)),
         state(StateSpace::Null()),
         scratch(StateSpace::Null()) {
-    bool denormals_are_zeros;
     is_valid = false;
     is_noisy = noisy;
 
@@ -617,12 +608,6 @@ class SimulatorHelper {
       StateSpace state_space = factory.CreateStateSpace();
       state = state_space.Create(num_qubits);
 
-      if (denormals_are_zeros) {
-        SetFlushToZeroAndDenormalsAreZeros();
-      } else {
-        ClearFlushToZeroAndDenormalsAreZeros();
-      }
-
       is_valid = true;
     } catch (const std::invalid_argument &exp) {
       // If this triggers, is_valid is false.
@@ -663,6 +648,7 @@ class SimulatorHelper {
 
   template <typename StateType>
   bool simulate(const StateType& input_state) {
+    ScopedFlushToZeroAndDenormalsAreZeros mxcsr_guard(denormals_are_zeros);
     init_state(input_state);
     bool result = false;
 
@@ -683,6 +669,7 @@ class SimulatorHelper {
   }
 
   bool simulate_subcircuit(uint64_t begin, uint64_t end) {
+    ScopedFlushToZeroAndDenormalsAreZeros mxcsr_guard(denormals_are_zeros);
     bool result = false;
 
     if (is_noisy) {
@@ -711,6 +698,7 @@ class SimulatorHelper {
   }
 
   std::vector<uint64_t> sample(uint64_t num_samples) {
+    ScopedFlushToZeroAndDenormalsAreZeros mxcsr_guard(denormals_are_zeros);
     StateSpace state_space = factory.CreateStateSpace();
     return state_space.Sample(state, num_samples, seed);
   }
@@ -727,6 +715,7 @@ class SimulatorHelper {
   std::vector<std::complex<double>> get_expectation_value(
       const std::vector<std::tuple<std::vector<OpString<float>>,
                                    unsigned>>& opsums_and_qubit_counts) {
+    ScopedFlushToZeroAndDenormalsAreZeros mxcsr_guard(denormals_are_zeros);
     Simulator simulator = factory.CreateSimulator();
     StateSpace state_space = factory.CreateStateSpace();
     using Fuser = MultiQubitGateFuser<IO>;
@@ -770,6 +759,7 @@ class SimulatorHelper {
   unsigned max_fused_size;
   unsigned verbosity;
   unsigned seed;
+  bool denormals_are_zeros;
 
   // Only set to "true" once initialization is complete.
   bool is_valid;
@@ -928,11 +918,7 @@ std::vector<unsigned> qsim_sample(const py::dict &options) {
     State state = state_space.Create(circuit.num_qubits);
     state_space.SetStateZero(state);
 
-    if (denormals_are_zeros) {
-      SetFlushToZeroAndDenormalsAreZeros();
-    } else {
-      ClearFlushToZeroAndDenormalsAreZeros();
-    }
+    ScopedFlushToZeroAndDenormalsAreZeros mxcsr_guard(denormals_are_zeros);
 
     if (!Runner::Run(param, factory, circuit, state, results)) {
       IO::errorf("qsim sampling of the circuit errored out.\n");
@@ -1008,11 +994,7 @@ std::vector<unsigned> qtrajectory_sample(const py::dict &options) {
       }
     };
 
-    if (denormals_are_zeros) {
-      SetFlushToZeroAndDenormalsAreZeros();
-    } else {
-      ClearFlushToZeroAndDenormalsAreZeros();
-    }
+    ScopedFlushToZeroAndDenormalsAreZeros mxcsr_guard(denormals_are_zeros);
 
     if (!NoisyRunner::RunBatch(param, ncircuit, seed, seed + 1,
                                state_space, simulator, measure)) {

qsimcirq_tests/denormals_are_zeros_test.py

@@ -0,0 +1,42 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import ctypes
+
+import cirq
+
+import qsimcirq
+
+
+def test_denormals_are_zeros_restored():
+    orig_value = ctypes.c_float(1e-40).value
+
+    _ = qsimcirq.QSimSimulator(
+        qsim_options=qsimcirq.QSimOptions(denormals_are_zeros=False)
+    ).simulate(cirq.Circuit())
+    assert (
+        ctypes.c_float(1e-40).value == orig_value
+    ), "denormals_are_zeros=False failed to restore flags"
+
+    _ = qsimcirq.QSimSimulator(
+        qsim_options=qsimcirq.QSimOptions(denormals_are_zeros=True)
+    ).simulate(cirq.Circuit())
+    assert (
+        ctypes.c_float(1e-40).value == orig_value
+    ), "denormals_are_zeros=True failed to restore flags"
+
+    _ = qsimcirq.QSimSimulator().simulate(cirq.Circuit())
+    assert (
+        ctypes.c_float(1e-40).value == orig_value
+    ), "default simulation failed to restore flags"