Replace Xoshiro with ChaCha in kickmix simulator

docs/kickmix/api_reference.md

@@ -4391,8 +4391,10 @@ def __init__(
             per-shot population counters on every executed instruction,
             slowing down fast gate simulation by ~2-4x and allocating
             extra counter memory, so it defaults to `False`.
-        seed: Optional PRNG seed for deterministic `HMR` and `R`
-            operations. If `None` (the default), seeds from OS entropy.
+        seed: Optional 256-bit integer PRNG seed (`0 <= seed < 2**256`)
+            for deterministic `HMR` and `R` operations (e.g. a SHA-256
+            digest `int(hashlib.sha256(...).hexdigest(), 16)`). If
+            `None` (the default), seeds from OS entropy.
 
     Examples:
         >>> import kickmix as km

kickmix/src/kickmix/py/interactive_simulator.pybind.cc

@@ -21,12 +21,31 @@ static std::mt19937_64 externally_seeded_rng() {
     return result;
 }
 
-static std::mt19937_64 make_rng(std::optional<uint64_t> seed) {
-    if (seed.has_value()) {
-        std::seed_seq seq{static_cast<uint32_t>(*seed), static_cast<uint32_t>(*seed >> 32)};
-        return std::mt19937_64(seq);
+static std::mt19937_64 make_rng(const pybind11::object &seed) {
+    if (seed.is_none()) {
+        return externally_seeded_rng();
     }
-    return externally_seeded_rng();
+    if (!pybind11::isinstance<pybind11::int_>(seed) || pybind11::isinstance<pybind11::bool_>(seed)) {
+        std::stringstream ss;
+        ss << "Expected seed to be an int or None, but got seed=" << pybind11::repr(seed);
+        throw pybind11::type_error(ss.str());
+    }
+    std::array<uint32_t, 8> seed_words{};
+    if (!seed.equal(pybind11::cast(0))) {
+        Py_ssize_t num_bytes = PyLong_AsNativeBytes(
+            seed.ptr(),
+            seed_words.data(),
+            sizeof(seed_words),
+            Py_ASNATIVEBYTES_LITTLE_ENDIAN | Py_ASNATIVEBYTES_REJECT_NEGATIVE | Py_ASNATIVEBYTES_UNSIGNED_BUFFER);
+        if (num_bytes < 0 || static_cast<size_t>(num_bytes) > sizeof(seed_words)) {
+            PyErr_Clear();
+            std::stringstream ss;
+            ss << "seed=" << pybind11::repr(seed) << " not in range(2**256)";
+            throw std::invalid_argument(ss.str());
+        }
+    }
+    std::seed_seq seq(seed_words.begin(), seed_words.end());
+    return std::mt19937_64(seq);
 }
 
 static void sim_clear_for_shot(PyInteractiveSimulator &self) {
@@ -114,16 +133,19 @@ pybind11::object int_obj_from_bytes_with_max_bits(uint8_t *bytes_buffer, size_t
 }
 
 PyInteractiveSimulator::PyInteractiveSimulator(
-    size_t init_batch_size,
-    pybind11::object init_debug_print,
-    bool init_counting_operations,
-    std::optional<uint64_t> seed) {
+    size_t init_batch_size, pybind11::object init_debug_print, bool init_counting_operations, pybind11::object seed) {
     batch_size = init_batch_size;
     counting_operations = init_counting_operations;
     std::mt19937_64 master_rng = make_rng(seed);
     with_sims([&](auto &sims) {
         for (size_t k = 0; k < batch_size; k += PY_SIM_WORD_BITS) {
-            std::seed_seq seq{master_rng(), master_rng(), master_rng(), master_rng()};
+            std::array<uint32_t, 8> chunk_seed{};
+            for (size_t w = 0; w < 4; w++) {
+                uint64_t r = master_rng();
+                chunk_seed[2 * w] = static_cast<uint32_t>(r);
+                chunk_seed[2 * w + 1] = static_cast<uint32_t>(r >> 32);
+            }
+            std::seed_seq seq(chunk_seed.begin(), chunk_seed.end());
             sims.emplace_back(std::mt19937_64(seq));
         }
         ensure_byte_buf_can_store_bits(sims.size() * PY_SIM_WORD_BITS);
@@ -778,9 +800,9 @@ void write_within_shot(
 void kickmix_py::register_interactive_simulator_methods(pybind11::class_<PyInteractiveSimulator> &c_sim) {
     c_sim.def(
         pybind11::init(
-            [](size_t batch_size, pybind11::object debug_print, bool counting_operations, std::optional<uint64_t> seed)
+            [](size_t batch_size, pybind11::object debug_print, bool counting_operations, pybind11::object seed)
                 -> PyInteractiveSimulator {
-                return PyInteractiveSimulator(batch_size, std::move(debug_print), counting_operations, seed);
+                return PyInteractiveSimulator(batch_size, std::move(debug_print), counting_operations, std::move(seed));
             }),
         pybind11::arg("batch_size"),
         pybind11::kw_only(),
@@ -808,8 +830,10 @@ void kickmix_py::register_interactive_simulator_methods(pybind11::class_<PyInter
                     per-shot population counters on every executed instruction,
                     slowing down fast gate simulation by ~2-4x and allocating
                     extra counter memory, so it defaults to `False`.
-                seed: Optional PRNG seed for deterministic `HMR` and `R`
-                    operations. If `None` (the default), seeds from OS entropy.
+                seed: Optional 256-bit integer PRNG seed (`0 <= seed < 2**256`)
+                    for deterministic `HMR` and `R` operations (e.g. a SHA-256
+                    digest `int(hashlib.sha256(...).hexdigest(), 16)`). If
+                    `None` (the default), seeds from OS entropy.
 
             Examples:
                 >>> import kickmix as km

kickmix/src/kickmix/py/interactive_simulator.pybind.h

@@ -35,7 +35,7 @@ struct PyInteractiveSimulator {
         size_t batch_size,
         pybind11::object debug_print = pybind11::bool_(true),
         bool counting_operations = false,
-        std::optional<uint64_t> seed = std::nullopt);
+        pybind11::object seed = pybind11::none());
     void ensure_big_enough_state_for(size_t num_qubits, size_t num_bits);
     void ensure_big_enough_state_for(kickmix::QubitOrBit val);
 };

kickmix/src/kickmix/py/interactive_simulator_test.py

@@ -620,15 +620,19 @@ def test_counting_operations_batch_sizes(batch_size: int):
 
 
 def test_simulator_seed():
+    import hashlib
+
     circuit = km.Circuit("""
         X q0
         X q1
         HMR q0 b0
         R q1
     """)
-    sim1 = km.Simulator(batch_size=512, seed=12345)
-    sim2 = km.Simulator(batch_size=512, seed=12345, counting_operations=True)
-    sim3 = km.Simulator(batch_size=512, seed=99999)
+    sha_seed = int(hashlib.sha256(str(circuit).encode()).hexdigest(), 16)
+    sim1 = km.Simulator(batch_size=512, seed=sha_seed)
+    sim2 = km.Simulator(batch_size=512, seed=sha_seed, counting_operations=True)
+    # Differ only in the highest bit (bit 255) to verify all 256 seed bits are used.
+    sim3 = km.Simulator(batch_size=512, seed=sha_seed ^ (1 << 255))
     sim1.do(circuit)
     sim2.do(circuit)
     sim3.do(circuit)
@@ -643,8 +647,21 @@ def test_simulator_seed():
     # Identical seeds produce identical HMR and R outcomes across all chunks.
     assert b0_1 == b0_2
     assert phase_1 == phase_2
-    # Different seeds produce different outcomes.
+    # Different seeds (even differing only in bit 255) produce different outcomes.
     assert b0_1 != b0_3
     assert phase_1 != phase_3
     # Independent 256-shot chunks within the same simulator do not repeat the same 256 bits.
     assert (b0_1 & ((1 << 256) - 1)) != (b0_1 >> 256)
+
+    # Boundary values 0 and 2**256 - 1 are accepted.
+    km.Simulator(batch_size=256, seed=0)
+    km.Simulator(batch_size=256, seed=(1 << 256) - 1)
+
+    with pytest.raises(ValueError, match="not in range"):
+        km.Simulator(batch_size=256, seed=-1)
+    with pytest.raises(ValueError, match="not in range"):
+        km.Simulator(batch_size=256, seed=1 << 256)
+    with pytest.raises(TypeError, match="Expected seed to be an int or None"):
+        km.Simulator(batch_size=256, seed=True)  # type: ignore[arg-type]
+    with pytest.raises(TypeError, match="Expected seed to be an int or None"):
+        km.Simulator(batch_size=256, seed="bad")  # type: ignore[arg-type]

kickmix/src/kickmix/sim/sim.h

@@ -10,9 +10,9 @@
 #include "kickmix/circuit/circuit.h"
 #include "kickmix/id/qubit_or_bit_or_bool.h"
 #include "kickmix/id/register_id.h"
+#include "kickmix/util/chacha.h"
 #include "kickmix/util/fixed_width_int.h"
 #include "kickmix/util/pop_counter.h"
-#include "kickmix/util/xoshiro_256_plusplus.h"
 
 namespace kickmix {
 
@@ -76,6 +76,7 @@ struct Sim {
     size_t num_qubits;
     size_t num_bits;
     std::mt19937_64 rng;
+    ChaCha8<TWord> chacha_rng;
     std::array<uint64_t, 64> transpose_buffer;
     std::vector<FixedPrecisionAngle128> angles;
     std::vector<RegisterData> registers;
@@ -105,30 +106,24 @@ struct Sim {
     }
 
     inline std::span<TWord> qubit_span() {
-        return {state_block.data() + 5, num_qubits};
+        return {state_block.data() + 1, num_qubits};
     }
     inline std::span<TWord> bit_span() {
-        return {state_block.data() + 5 + num_qubits, num_bits};
+        return {state_block.data() + 1 + num_qubits, num_bits};
     }
     inline TWord &global_phase_ref() {
         return state_block[0];
     }
-    inline std::span<TWord> rng_state_span() {
-        return {&state_block[1], 4};
-    }
 
     inline std::span<const TWord> qubit_span() const {
-        return {state_block.data() + 5, num_qubits};
+        return {state_block.data() + 1, num_qubits};
     }
     inline std::span<const TWord> bit_span() const {
-        return {state_block.data() + 5 + num_qubits, num_bits};
+        return {state_block.data() + 1 + num_qubits, num_bits};
     }
     inline const TWord &global_phase_ref() const {
         return state_block[0];
     }
-    inline std::span<const TWord> rng_state_span() const {
-        return {&state_block[1], 4};
-    }
 
     FixedPrecisionAngle128 read_shot_phase(size_t shot_index) const {
         auto result = angles[shot_index];
@@ -139,11 +134,8 @@ struct Sim {
     }
 
     explicit Sim(std::mt19937_64 &&rng)
-        : state_block(5), num_qubits(0), num_bits(0), rng(rng), transpose_buffer{}, angles(BATCH_SIZE) {
-        state_block[1].randomize(rng);
-        state_block[2].randomize(rng);
-        state_block[3].randomize(rng);
-        state_block[4].randomize(rng);
+        : state_block(1), num_qubits(0), num_bits(0), rng(rng), chacha_rng{}, transpose_buffer{}, angles(BATCH_SIZE) {
+        chacha_rng.randomize(this->rng);
     }
     Sim(const Sim &sim) = default;
     Sim(Sim &&sim) = default;
@@ -259,13 +251,13 @@ struct Sim {
         std::span<TWord> old_qubits = qubit_span();
         std::vector<TWord> old_state = std::move(state_block);
         state_block = std::vector<TWord>();
-        state_block.resize(min_num_qubits + min_num_bits + 5);
+        state_block.resize(min_num_qubits + min_num_bits + 1);
         size_t old_num_qubits = num_qubits;
         size_t old_num_bits = num_bits;
         num_qubits = min_num_qubits;
         num_bits = min_num_bits;
 
-        memcpy(state_block.data(), old_state.data(), 5 * sizeof(TWord));
+        memcpy(state_block.data(), old_state.data(), 1 * sizeof(TWord));
         if (old_num_bits) {
             memcpy(bit_span().data(), old_bits.data(), old_num_bits * sizeof(TWord));
         }
@@ -286,13 +278,12 @@ struct Sim {
 
         num_qubits = circuit.num_qubits;
         num_bits = circuit.num_bits;
-        state_block.resize(num_qubits + num_bits + 5);
+        state_block.resize(num_qubits + num_bits + 1);
         clear_for_shot();
     }
 
     void apply(const Circuit &circuit) {
-        Xoshiro256PlusPlus<TWord> rng_state;
-        memcpy(&rng_state, rng_state_span().data(), sizeof(rng_state));
+        auto &rng_state = chacha_rng;
         const auto qubits = qubit_span();
         const auto bits = bit_span();
         TWord current_base_condition = condition_stack.empty() ? ~TWord{} : condition_stack.back();
@@ -506,7 +497,6 @@ struct Sim {
         }
 
         global_phase_ref() = phase;
-        memcpy(rng_state_span().data(), &rng_state, sizeof(rng_state));
     }
 
     void clear_operation_statistics() {

kickmix/src/kickmix/sim/sim.test.cc

@@ -219,7 +219,7 @@ TEST(sim, safe_read_val_bounds) {
     Sim<b64, false> sim(INDEPENDENT_TEST_RNG());
     sim.num_qubits = 2;
     sim.num_bits = 2;
-    sim.state_block.assign(5 + 4, ~b64{});
+    sim.state_block.assign(1 + 4, ~b64{});
     ASSERT_EQ(sim.safe_read_val(QubitId(0)), ~b64{});
     ASSERT_EQ(sim.safe_read_val(QubitId(1)), ~b64{});
     ASSERT_EQ(sim.safe_read_val(QubitId(2)), b64{});

kickmix/src/kickmix/simd/b128_polyfill.h

@@ -103,6 +103,34 @@ struct alignas(sizeof(uint64_t) * 2) b128_polyfill {
     inline b128_polyfill operator|(const b128_polyfill &other) const {
         return {v[0] | other.v[0], v[1] | other.v[1]};
     }
+    inline b128_polyfill u32_add(const b128_polyfill &other) const {
+        const uint32_t *v1 = (const uint32_t *)&v;
+        const uint32_t *v2 = (const uint32_t *)&other.v;
+        b128_polyfill result;
+        uint32_t *v3 = (uint32_t *)&result.v;
+        for (size_t k = 0; k < BIT_SIZE / 32; k++) {
+            v3[k] = v1[k] + v2[k];
+        }
+        return result;
+    }
+    inline b128_polyfill u32_right_shift(size_t shift) const {
+        const uint32_t *v1 = (const uint32_t *)&v;
+        b128_polyfill result;
+        uint32_t *v3 = (uint32_t *)&result.v;
+        for (size_t k = 0; k < BIT_SIZE / 32; k++) {
+            v3[k] = v1[k] >> shift;
+        }
+        return result;
+    }
+    inline b128_polyfill u32_left_shift(size_t shift) const {
+        const uint32_t *v1 = (const uint32_t *)&v;
+        b128_polyfill result;
+        uint32_t *v3 = (uint32_t *)&result.v;
+        for (size_t k = 0; k < BIT_SIZE / 32; k++) {
+            v3[k] = v1[k] << shift;
+        }
+        return result;
+    }
     inline b128_polyfill u32_eq(const b128_polyfill &other) const {
         const uint32_t *v1 = (uint32_t *)&v;
         const uint32_t *v2 = (uint32_t *)&other.v;

kickmix/src/kickmix/simd/b128_sse.h

@@ -101,6 +101,15 @@ struct alignas(sizeof(uint64_t) * 2) b128_sse {
     inline b128_sse u64_left_shift(size_t shift) const {
         return {_mm_slli_epi64(v, shift)};
     }
+    inline b128_sse u32_add(const b128_sse &other) const {
+        return {_mm_add_epi32(v, other.v)};
+    }
+    inline b128_sse u32_right_shift(size_t shift) const {
+        return {_mm_srli_epi32(v, shift)};
+    }
+    inline b128_sse u32_left_shift(size_t shift) const {
+        return {_mm_slli_epi32(v, shift)};
+    }
     inline b128_sse u32_eq(const b128_sse &other) const {
         return {_mm_cmpeq_epi32(v, other.v)};
     }

kickmix/src/kickmix/simd/b256_avx.h

@@ -96,6 +96,15 @@ struct alignas(sizeof(uint64_t) * 4) b256_avx {
     inline b256_avx u64_left_shift(size_t shift) const {
         return {_mm256_slli_epi64(v, shift)};
     }
+    inline b256_avx u32_add(const b256_avx &other) const {
+        return {_mm256_add_epi32(v, other.v)};
+    }
+    inline b256_avx u32_right_shift(size_t shift) const {
+        return {_mm256_srli_epi32(v, shift)};
+    }
+    inline b256_avx u32_left_shift(size_t shift) const {
+        return {_mm256_slli_epi32(v, shift)};
+    }
     inline b256_avx u32_eq(const b256_avx &other) const {
         return {_mm256_cmpeq_epi32(v, other.v)};
     }

kickmix/src/kickmix/simd/b256_polyfill.h

@@ -134,6 +134,34 @@ struct alignas(sizeof(uint64_t) * 4) b256_polyfill {
             v[3] | other.v[3],
         };
     }
+    inline b256_polyfill u32_add(const b256_polyfill &other) const {
+        const uint32_t *v1 = (const uint32_t *)&v;
+        const uint32_t *v2 = (const uint32_t *)&other.v;
+        b256_polyfill result;
+        uint32_t *v3 = (uint32_t *)&result.v;
+        for (size_t k = 0; k < BIT_SIZE / 32; k++) {
+            v3[k] = v1[k] + v2[k];
+        }
+        return result;
+    }
+    inline b256_polyfill u32_right_shift(size_t shift) const {
+        const uint32_t *v1 = (const uint32_t *)&v;
+        b256_polyfill result;
+        uint32_t *v3 = (uint32_t *)&result.v;
+        for (size_t k = 0; k < BIT_SIZE / 32; k++) {
+            v3[k] = v1[k] >> shift;
+        }
+        return result;
+    }
+    inline b256_polyfill u32_left_shift(size_t shift) const {
+        const uint32_t *v1 = (const uint32_t *)&v;
+        b256_polyfill result;
+        uint32_t *v3 = (uint32_t *)&result.v;
+        for (size_t k = 0; k < BIT_SIZE / 32; k++) {
+            v3[k] = v1[k] << shift;
+        }
+        return result;
+    }
     inline b256_polyfill u32_eq(const b256_polyfill &other) const {
         const uint32_t *v1 = (uint32_t *)&v;
         const uint32_t *v2 = (uint32_t *)&other.v;

kickmix/src/kickmix/simd/b512_avx.h

@@ -97,6 +97,15 @@ struct alignas(sizeof(uint64_t) * 8) b512_avx {
     inline b512_avx u64_left_shift(size_t shift) const {
         return {_mm512_slli_epi64(v, shift)};
     }
+    inline b512_avx u32_add(const b512_avx &other) const {
+        return {_mm512_add_epi32(v, other.v)};
+    }
+    inline b512_avx u32_right_shift(size_t shift) const {
+        return {_mm512_srli_epi32(v, shift)};
+    }
+    inline b512_avx u32_left_shift(size_t shift) const {
+        return {_mm512_slli_epi32(v, shift)};
+    }
     inline b512_avx u32_eq(const b512_avx &other) const {
         return {_mm512_movm_epi32(_mm512_cmpeq_epi32_mask(v, other.v))};
     }