PennyLaneAI / catalyst

A JIT compiler for hybrid quantum programs in PennyLane
https://docs.pennylane.ai/projects/catalyst
Apache License 2.0
122 stars 27 forks source link

Add seeded execution to the Catalyst runtime #936

Closed paul0403 closed 1 month ago

paul0403 commented 1 month ago

Context: The test test_dynamic_one_shot_several_mcms in frontend/test/pytest/test_mid_circuit_measurement.py was marked as skipped because it was flaky: #842

After some investigation, it was found out that the test involves a measurement, which gives random results. The qjit run of the test is random, but the reference default.qubit run is seeded, so sometimes the results fall outside the acceptable range of np.allclose.

To resolve this, we add a seeding mechanism for qjit.

Description of the Change: Implemented a random seeding infrastructure for qjit.

The top-level qjit decorator can now take in a string argument seed="some_string". The default value is empty string, which means an unseeded run.

The string will be propagated to the runtime ExecutionContext, which then initializes a PRNG (the c++ std::mt19937 pseudo random number generator) in the context. The seed and the PRNG is then sent to individual devices. When performing measurements, the devices draw according to this PRNG.

A proper seed is found and set, to deterministically resolve the flaky test test_dynamic_one_shot_several_mcms.

Benefits:

Possible Drawbacks:

Related GitHub Issues: closes #839

[sc-66696]

codecov[bot] commented 1 month ago

Codecov Report

Attention: Patch coverage is 98.11321% with 2 lines in your changes missing coverage. Please review.

Project coverage is 97.90%. Comparing base (897630c) to head (276761e).

Files Patch % Lines
runtime/include/QuantumDevice.hpp 0.00% 1 Missing :warning:
runtime/lib/capi/ExecutionContext.hpp 75.00% 1 Missing :warning:
Additional details and impacted files ```diff @@ Coverage Diff @@ ## main #936 +/- ## ======================================= Coverage 97.90% 97.90% ======================================= Files 72 73 +1 Lines 10341 10371 +30 Branches 1180 1185 +5 ======================================= + Hits 10124 10154 +30 - Misses 171 172 +1 + Partials 46 45 -1 ```

:umbrella: View full report in Codecov by Sentry.
:loudspeaker: Have feedback on the report? Share it here.

paul0403 commented 1 month ago

Confirming that this is what we want (using "draw" as a name for the internal "random state" in lightning simulator runtime) when we set a seed:

i.e.

# in file test.py

dev = qml.device("lightning.qubit", wires=1, shots=None)
@qjit(seed="qwerty")
@qml.qnode(dev)
def circuit():
    qml.Hadamard(0)
    m = measure(0)
    @cond(m)
    def cfun0():
        qml.Hadamard(0)
    cfun0()
    return qml.probs()

print(circuit(), circuit(), circuit(), circuit())

If we run python test.py, the 4 executions should give "random" results from each other; however, running python test.py again should produce the same 4 results:

$ python3 test.py
[0.5 0.5] [0.5 0.5] [1. 0.] [0.5 0.5]
draw is 0.73658 draw is 0.797645 draw is 0.267068 draw is 0.793337 
$ python3 test.py
[0.5 0.5] [0.5 0.5] [1. 0.] [0.5 0.5]
draw is 0.73658 draw is 0.797645 draw is 0.267068 draw is 0.793337

The above is exact (shots=None). If we use a high number of shots (e,g, shots=10000) the probabilities will have small fluctuations:

$ python3 flaky_mcm.py
[0.5013 0.4987] [1. 0.] [1. 0.] [1. 0.]
draw is 0.838864 draw is 0.289431 draw is 0.405051 draw is 0.0953809 

$ python3 flaky_mcm.py
[0.5021 0.4979] [1. 0.] [1. 0.] [1. 0.]
draw is 0.838864 draw is 0.289431 draw is 0.405051 draw is 0.0953809

$ python3 flaky_mcm.py
[0.5019 0.4981] [1. 0.] [1. 0.] [1. 0.]
draw is 0.838864 draw is 0.289431 draw is 0.405051 draw is 0.0953809

$ python3 flaky_mcm.py
[0.4899 0.5101] [1. 0.] [1. 0.] [1. 0.]
draw is 0.838864 draw is 0.289431 draw is 0.405051 draw is 0.0953809

@dime10 @mudit2812

dime10 commented 1 month ago

@paul0403 No, I would say each execution of the qjit function should produce the same results, since it is seeded.

However, if you were to invoke the same QNode multiple times within a single execution of the qjit function, they should produce different results.

paul0403 commented 1 month ago

@dime10 Ah, so

dev = qml.device("lightning.qubit", wires=1, shots=None)

@qjit(seed="qwerty")
def workflow():
    @qml.qnode(dev)
    def circuit():
        qml.Hadamard(0)
        m = measure(0)
        @cond(m)
        def cfun0():
            qml.Hadamard(0)
        cfun0()
        return qml.probs()

    return circuit(), circuit(), circuit(), circuit()

print(workflow())
print("#######")
print(workflow())
(Array([0.4975, 0.5025], dtype=float64), Array([1., 0.], dtype=float64), Array([1., 0.], dtype=float64), Array([1., 0.], dtype=float64))
#######
(Array([0.4975, 0.5025], dtype=float64), Array([1., 0.], dtype=float64), Array([1., 0.], dtype=float64), Array([1., 0.], dtype=float64))

i.e. the history of both workflow (which is the QJIT) calls are the same, but within each QJIT call the qnode results are different?

dime10 commented 1 month ago

@dime10 Ah, so the history of both workflow (which is the QJIT) calls are the same, but within each QJIT call the qnode results are different?

Exactly :)

mudit2812 commented 1 month ago

@dime10 @paul0403 When we talked yesterday the behaviour I had in mind was more consistent with the first implementation. This is really more of a product question though.

paul0403 commented 1 month ago

@dime10 @paul0403 When we talked yesterday the behaviour I had in mind was more consistent with the first implementation. This is really more of a product question though.

@josh146

dime10 commented 1 month ago

@dime10 @paul0403 When we talked yesterday the behaviour I had in mind was more consistent with the first implementation. This is really more of a product question though.

That would be harder to realize because there is no persistent state in the runtime across qjit invocations. (and also unnecessary imo)

paul0403 commented 1 month ago

~TODO: add frontend tests for seeded qjit; changelog~ done

paul0403 commented 1 month ago

Question: right now only measurements are seeded. Should sampling (aka shots) be seeded as well?

paul0403 commented 1 month ago

@erick-xanadu @dime10 Thanks for the suggestions! Yes, the root of the complexity is C not allowing overloading.

If we are fine with changing the signature to __catalyst__rt__initilize(char *) everywhere (and modifying all the necessary tests to reflect this), then this is by far the cleanest approach imo. The complex blocks in the conversion pass also become unnecessary.

I have done this in the most recent commit ~(without going over the tests right now)~.

paul0403 commented 1 month ago

Codecov complains about coverage on the prng getter and setter in the QuantumDevice base class. Three tests are added in the runtime currently to cover this:

  1. Test_LightningDriver.cpp: test that the prng g(s)etter will g(s)et correctly on lightning and kokkos
  2. Test_OpenQasmDevice.cpp: test that the prng g(s)etter correctly does nothing on openqasm (since it's the default definition)
  3. Test_LightningMeasures.cpp: uses the setter to set a seed for a lightning/kokkos device, and then measure. Test that the measurement results are consistent on two devices with the same seed.

I think these tests are enough coverage in the runtime.