Add seeded execution to the Catalyst runtime

frontend/catalyst/compiler.py

@@ -84,6 +84,7 @@ class CompileOptions:
     static_argnums: Optional[Union[int, Iterable[int]]] = None
     abstracted_axes: Optional[Union[Iterable[Iterable[str]], Dict[int, str]]] = None
     lower_to_llvm: Optional[bool] = True
+    seed: Optional[str] = ""
 
     def __post_init__(self):
         # Make the format of static_argnums easier to handle.

frontend/catalyst/jit.py

@@ -80,6 +80,7 @@ def qjit(
     pipelines=None,
     static_argnums=None,
     abstracted_axes=None,
+    seed="",
 ):  # pylint: disable=too-many-arguments,unused-argument
     """A just-in-time decorator for PennyLane and JAX programs using Catalyst.
 
@@ -123,6 +124,11 @@ def qjit(
             Function arguments with ``abstracted_axes`` specified will be compiled to ranked tensors
             with dynamic shapes. For more details, please see the Dynamically-shaped Arrays section
             below.
+        seed (str):
+            The seed for random operations in a qjit call, such as circuit measurement results.
+            The default value is an empty string, which means no seeding is performed, and all
+            processes are random.
+            Note that if the circuit is run from shots, the sampled results are NOT seeded.
 
     Returns:
         QJIT object.
@@ -652,7 +658,7 @@ def generate_ir(self):
         mlir_module, ctx = lower_jaxpr_to_mlir(self.jaxpr, self.__name__)
 
         # Inject Runtime Library-specific functions (e.g. setup/teardown).
-        inject_functions(mlir_module, ctx)
+        inject_functions(mlir_module, ctx, self.compile_options.seed)
 
         # Canonicalize the MLIR since there can be a lot of redundancy coming from JAX.
         options = copy.deepcopy(self.compile_options)

frontend/catalyst/third_party/oqc/src/OQCDevice.cpp

@@ -55,6 +55,10 @@ void OQCDevice::SetDeviceShots([[maybe_unused]] size_t shots) { device_shots = s
 
 auto OQCDevice::GetDeviceShots() const -> size_t { return device_shots; }
 
+void OQCDevice::SetDeviceSeed([[maybe_unused]] std::string _seed) { seed = _seed; }
+
+void OQCDevice::SetDevicePRNG([[maybe_unused]] std::mt19937 *_gen) { gen = _gen; }
+
 auto OQCDevice::Zero() const -> Result { return const_cast<Result>(&GLOBAL_RESULT_FALSE_CONST); }
 
 auto OQCDevice::One() const -> Result { return const_cast<Result>(&GLOBAL_RESULT_TRUE_CONST); }

frontend/catalyst/third_party/oqc/src/OQCDevice.hpp

@@ -52,6 +52,9 @@ class OQCDevice final : public Catalyst::Runtime::QuantumDevice {
     Catalyst::Runtime::CacheManager<std::complex<double>> cache_manager{};
     bool tape_recording{false};
     size_t device_shots;
+
+    std::string seed;
+    std::mt19937 *gen;
 
     std::unordered_map<std::string, std::string> device_kwargs;
 

frontend/catalyst/utils/gen_mlir.py

@@ -20,13 +20,21 @@
 from jaxlib.mlir.dialects._func_ops_gen import FuncOp
 
 
-def gen_setup(ctx):
+def gen_setup(ctx, seed):
     """
     This function returns an MLIR module with the "setup" function. The setup
     function is a function that needs to be called before calling a
     JIT-compiled function. It initializes the global device context in the runtime.
     """
-    txt = """
+    if seed != "":
+        txt = f"""
+func.func @setup() -> () {{
+    "quantum.init"() {{seed = "{seed}"}} : () -> ()
+    return
+}}
+"""
+    else:
+        txt = """
 func.func @setup() -> () {
     "quantum.init"() : () -> ()
     return
@@ -50,14 +58,14 @@ def gen_teardown(ctx):
     return ir.Module.parse(txt, ctx)
 
 
-def inject_functions(module, ctx):
+def inject_functions(module, ctx, seed):
     """
     This function appends functions to the input module.
     """
     # Add C interface for the quantum function.
     module.body.operations[0].attributes["llvm.emit_c_interface"] = ir.UnitAttr.get(context=ctx)
 
-    setup_module = gen_setup(ctx)
+    setup_module = gen_setup(ctx, seed)
     setup_func = setup_module.body.operations[0]
     module.body.append(setup_func)
 

frontend/test/pytest/test_mid_circuit_measurement.py

@@ -496,11 +496,6 @@ def outer():
         assert result.shape == (shots,)
         assert jnp.allclose(result, 1.0)
 
-    # TODO: dynamic_one_shot_several_mcms is a flaky test.
-    # We remove this test for now and revisit in the future.
-    @pytest.mark.skip(
-        reason="dynamic_one_shot_several_mcms is a flaky test and needs further investigation"
-    )
     @pytest.mark.parametrize("shots", [10000])
     @pytest.mark.parametrize("postselect", [None, 0, 1])
     @pytest.mark.parametrize("measure_f", [qml.counts, qml.expval, qml.probs, qml.sample, qml.var])
@@ -542,7 +537,7 @@ def ref_func(x, y):
 
         dev = qml.device(backend, wires=2, shots=shots)
 
-        @qjit
+        @qjit(seed="8v4Lj7L")
         @qml.qnode(dev, postselect_mode=postselect_mode, mcm_method="one-shot")
         def func(x, y):
             qml.RX(x, 0)

mlir/lib/Quantum/Transforms/ConversionPatterns.cpp

@@ -154,17 +154,36 @@ template <typename T> struct RTBasedPattern : public OpConversionPattern<T> {
         StringRef qirName;
         if constexpr (std::is_same_v<T, InitializeOp>) {
             qirName = "__catalyst__rt__initialize";
+            InitializeOp InitOp = cast<InitializeOp>(op);
+            Location loc = InitOp.getLoc();
+            ModuleOp mod = InitOp->getParentOfType<ModuleOp>();
+            Type charPtrType = LLVM::LLVMPointerType::get(rewriter.getContext());
+            Type qirSignature = LLVM::LLVMFunctionType::get(LLVM::LLVMVoidType::get(ctx),
+                                                            /* seed = */ {charPtrType});
+            Value seed_gs;
+            if (InitOp->hasAttr("seed")) {
+                auto seed_str = cast<StringAttr>(InitOp->getAttr("seed")).str();
+                seed_gs = getGlobalString(loc, rewriter, seed_str,
+                                          StringRef(seed_str.c_str(), seed_str.length() + 1), mod);
+            }
+            else {
+                seed_gs = getGlobalString(loc, rewriter, "unseeded",
+                                          "__catalyst__unseeded__qjit__run__", mod);
+            }
+            LLVM::LLVMFuncOp fnDecl =
+                ensureFunctionDeclaration(rewriter, op, qirName, qirSignature);
+            SmallVector<Value> operands = {seed_gs};
+            rewriter.create<LLVM::CallOp>(loc, fnDecl, operands);
+            rewriter.eraseOp(op);
         }
         else {
             qirName = "__catalyst__rt__finalize";
+            Type qirSignature = LLVM::LLVMFunctionType::get(LLVM::LLVMVoidType::get(ctx), {});
+            LLVM::LLVMFuncOp fnDecl =
+                ensureFunctionDeclaration(rewriter, op, qirName, qirSignature);
+            rewriter.replaceOpWithNewOp<LLVM::CallOp>(op, fnDecl, ValueRange{});
         }
 
-        Type qirSignature = LLVM::LLVMFunctionType::get(LLVM::LLVMVoidType::get(ctx), {});
-
-        LLVM::LLVMFuncOp fnDecl = ensureFunctionDeclaration(rewriter, op, qirName, qirSignature);
-
-        rewriter.replaceOpWithNewOp<LLVM::CallOp>(op, fnDecl, ValueRange{});
-
         return success();
     }
 };

mlir/test/Quantum/ConversionTest.mlir

@@ -18,12 +18,12 @@
 // Runtime Management //
 ////////////////////////
 
-// CHECK: llvm.func @__catalyst__rt__initialize()
+// CHECK: llvm.func @__catalyst__rt__initialize(!llvm.ptr)
 
 // CHECK-LABEL: @init
 func.func @init() {
 
-    // CHECK: llvm.call @__catalyst__rt__initialize()
+    // CHECK: llvm.call @__catalyst__rt__initialize({{%.+}}) : (!llvm.ptr) -> ()
     quantum.init
 
     return

runtime/include/QuantumDevice.hpp

@@ -17,6 +17,7 @@
 #include <complex>
 #include <memory>
 #include <optional>
+#include <random>
 #include <vector>
 
 #include "DataView.hpp"
@@ -103,6 +104,20 @@ struct QuantumDevice {
      */
     [[nodiscard]] virtual auto GetDeviceShots() const -> size_t = 0;
 
+    /**
+     * @brief Set the PRNG seed of the device.
+     *
+     * @param seed The PRNG seed.
+     */
+    virtual void SetDeviceSeed(std::string seed) = 0;
+
+    /**
+     * @brief Set the PRNG of the device.
+     *
+     * @param gen The std::mt19937 PRNG object.
+     */
+    virtual void SetDevicePRNG(std::mt19937 *gen) = 0;
+
     /**
      * @brief Start recording a quantum tape if provided.
      *

runtime/include/RuntimeCAPI.h

@@ -24,7 +24,7 @@ extern "C" {
 
 // Quantum Runtime Instructions
 void __catalyst__rt__fail_cstr(const char *);
-void __catalyst__rt__initialize();
+void __catalyst__rt__initialize(char *seed = nullptr);
 void __catalyst__rt__device_init(int8_t *, int8_t *, int8_t *);
 void __catalyst__rt__device_release();
 void __catalyst__rt__finalize();

runtime/lib/backend/common/Utils.hpp

@@ -44,6 +44,8 @@
     void StopTapeRecording() override;                                                             \
     void SetDeviceShots(size_t shots) override;                                                    \
     [[nodiscard]] auto GetDeviceShots() const->size_t override;                                    \
+    void SetDeviceSeed(std::string seed) override;                                                 \
+    void SetDevicePRNG(std::mt19937 *gen) override;                                                \
     void PrintState() override;                                                                    \
     [[nodiscard]] auto Zero() const->Result override;                                              \
     [[nodiscard]] auto One() const->Result override;
@@ -271,8 +273,8 @@ constexpr auto has_gate(const SimulatorGateInfoDataT<size> &arr, const std::stri
     return false;
 }
 
-static inline auto simulateDraw(const std::vector<double> &probs, std::optional<int32_t> postselect)
-    -> bool
+static inline auto simulateDraw(const std::vector<double> &probs, std::optional<int32_t> postselect,
+                                std::mt19937 *gen, bool has_seed) -> bool
 {
     if (postselect) {
         auto postselect_value = postselect.value();
@@ -283,10 +285,19 @@ static inline auto simulateDraw(const std::vector<double> &probs, std::optional<
 
     // Normal flow, no post-selection
     // Draw a number according to the given distribution
-    std::random_device rd;
-    std::mt19937 gen(rd());
     std::uniform_real_distribution<> dis(0., 1.);
-    float draw = dis(gen);
+
+    float draw;
+    if (has_seed) {
+        draw = dis(*gen);
+        (*gen)();
+    }
+    else {
+        std::random_device rd;
+        std::mt19937 gen_no_seed(rd());
+        draw = dis(gen_no_seed);
+    }
+
     return draw > probs[0];
 }
 

runtime/lib/backend/dummy/dummy_device.cpp

@@ -39,6 +39,8 @@ struct DummyDevice final : public Catalyst::Runtime::QuantumDevice {
     [[nodiscard]] auto GetNumQubits() const -> size_t override { return 0; }
     void SetDeviceShots(size_t shots) override {}
     [[nodiscard]] auto GetDeviceShots() const -> size_t override { return 0; }
+    void SetDeviceSeed(std::string seed) override {}
+    void SetDevicePRNG(std::mt19937 *gen) override {}
     void StartTapeRecording() override {}
     void StopTapeRecording() override {}
     void PrintState() override {}

runtime/lib/backend/lightning/lightning_dynamic/LightningSimulator.cpp

@@ -85,6 +85,10 @@ void LightningSimulator::SetDeviceShots(size_t shots) { this->device_shots = sho
 
 auto LightningSimulator::GetDeviceShots() const -> size_t { return this->device_shots; }
 
+void LightningSimulator::SetDeviceSeed(std::string _seed) { this->seed = _seed; }
+
+void LightningSimulator::SetDevicePRNG(std::mt19937 *_gen) { this->gen = _gen; }
+
 void LightningSimulator::PrintState()
 {
     using std::cout;
@@ -433,7 +437,7 @@ auto LightningSimulator::Measure(QubitIdType wire, std::optional<int32_t> postse
     SetDeviceShots(device_shots);
 
     // It represents the measured result, true for 1, false for 0
-    bool mres = Lightning::simulateDraw(probs, postselect);
+    bool mres = Lightning::simulateDraw(probs, postselect, this->gen, this->hasSeed());
     auto dev_wires = getDeviceWires(wires);
     this->device_sv->collapse(dev_wires[0], mres ? 1 : 0);
     return mres ? this->One() : this->Zero();

runtime/lib/backend/lightning/lightning_dynamic/LightningSimulator.hpp

@@ -51,6 +51,9 @@ class LightningSimulator final : public Catalyst::Runtime::QuantumDevice {
     bool tape_recording{false};
     size_t device_shots;
 
+    std::string seed;
+    std::mt19937 *gen;
+
     bool mcmc{false};
     size_t num_burnin{0};
     std::string kernel_name;
@@ -84,6 +87,8 @@ class LightningSimulator final : public Catalyst::Runtime::QuantumDevice {
         return res;
     }
 
+    inline auto hasSeed() -> bool { return this->seed != ""; }
+
   public:
     explicit LightningSimulator(const std::string &kwargs = "{}")
     {

runtime/lib/backend/lightning/lightning_kokkos/LightningKokkosSimulator.cpp

@@ -100,6 +100,10 @@ void LightningKokkosSimulator::SetDeviceShots(size_t shots) { this->device_shots
 
 auto LightningKokkosSimulator::GetDeviceShots() const -> size_t { return this->device_shots; }
 
+void LightningKokkosSimulator::SetDeviceSeed(std::string _seed) { this->seed = _seed; }
+
+void LightningKokkosSimulator::SetDevicePRNG(std::mt19937 *_gen) { this->gen = _gen; }
+
 void LightningKokkosSimulator::PrintState()
 {
     using std::cout;
@@ -455,7 +459,7 @@ auto LightningKokkosSimulator::Measure(QubitIdType wire, std::optional<int32_t>
     SetDeviceShots(device_shots);
 
     // It represents the measured result, true for 1, false for 0
-    bool mres = Lightning::simulateDraw(probs, postselect);
+    bool mres = Lightning::simulateDraw(probs, postselect, this->gen, this->hasSeed());
     auto dev_wires = getDeviceWires(wires);
     this->device_sv->collapse(dev_wires[0], mres ? 1 : 0);
     return mres ? this->One() : this->Zero();

runtime/lib/backend/lightning/lightning_kokkos/LightningKokkosSimulator.hpp

@@ -49,6 +49,9 @@ class LightningKokkosSimulator final : public Catalyst::Runtime::QuantumDevice {
 
     size_t device_shots;
 
+    std::string seed;
+    std::mt19937 *gen;
+
     std::unique_ptr<StateVectorT> device_sv = std::make_unique<StateVectorT>(0);
     LightningKokkosObsManager<double> obs_manager{};
 
@@ -78,6 +81,8 @@ class LightningKokkosSimulator final : public Catalyst::Runtime::QuantumDevice {
         return res;
     }
 
+    inline auto hasSeed() -> bool { return this->seed != ""; }
+
   public:
     explicit LightningKokkosSimulator(const std::string &kwargs = "{}")
     {

runtime/lib/backend/openqasm/OpenQasmDevice.cpp

@@ -80,6 +80,10 @@ void OpenQasmDevice::SetDeviceShots([[maybe_unused]] size_t shots) { device_shot
 
 auto OpenQasmDevice::GetDeviceShots() const -> size_t { return device_shots; }
 
+void OpenQasmDevice::SetDeviceSeed([[maybe_unused]] std::string _seed) { seed = _seed; }
+
+void OpenQasmDevice::SetDevicePRNG([[maybe_unused]] std::mt19937 *_gen) { gen = _gen; }
+
 void OpenQasmDevice::PrintState()
 {
     using std::cout;

runtime/lib/backend/openqasm/OpenQasmDevice.hpp

@@ -52,6 +52,9 @@ class OpenQasmDevice final : public Catalyst::Runtime::QuantumDevice {
     bool tape_recording{false};
     size_t device_shots;
 
+    std::string seed;
+    std::mt19937 *gen;
+
     OpenQasm::OpenQasmObsManager obs_manager{};
     OpenQasm::BuilderType builder_type;
     std::unordered_map<std::string, std::string> device_kwargs;