Correctly querying wire IDs for `SetState` and `SetBasisState` in runtime #1047

paul0403 · 2024-08-23T14:55:46Z

Context:
Recycling a device when a previous execution involves stateprep causes a crash. See #1044

Description of the Change:
In LightningSimulator, SetState and SetBasisState now correctly query the DevQubits from the SimQubits in its qubit_manger's map.

Benefits: We can now have multiple qnode functions involving stateprep in a workflow

Related GitHub Issues: closes #1044

[sc-71894]

…hether there is stateprep in a previous device to judge whether the device should be recycled.

dime10 · 2024-08-23T15:36:27Z

Should we not look into why it crashes in the first place?

paul0403 · 2024-08-23T17:40:08Z

Should we not look into why it crashes in the first place?

Yes, draft 1 is trying to contain this issue within Catalyst but I agree the root cause needs to be addressed in Lightning. Chatted with @maliasadi regarding how the two repos see each other (also tagging to keep Ali in the loop on this issue).

paul0403 · 2024-08-23T20:35:34Z

@erick-xanadu Root cause: not properly resetting the wires' va_list in __catalyst__qis__SetState:

catalyst/runtime/lib/capi/RuntimeCAPI.cpp

Line 467 in 1950dd1

wires[i] = va_arg(args, QubitIdType);

e.g. for a 2-wire circuit, the first execution will make this wires vector {0,1}, but the second execution makes it {2,3}, so causing "index out of range" style segfaults downstream in the backend, specifically https://github.com/PennyLaneAI/pennylane-lightning/blob/97554407e114db41690b665c01f342c8ec27b98a/pennylane_lightning/core/src/simulators/lightning_qubit/StateVectorLQubit.hpp#L809, where the iterator to std::vector::erase(it) is out of bounds.

Potentially this is not a va_list issue, but an issue with the device init/reset and allocation

Good news for @maliasadi @mlxd : from what I can tell this is not an issue with how setstate is implemented in lightning, but with how setstate is called in catalyst, aka no action required from lightning's end

paul0403 · 2024-08-23T21:33:10Z

A supernatural feeling when it's Friday 6pm and you spent a whole day on a bug and the fix is one line

runtime/lib/backend/common/QubitManager.hpp

codecov · 2024-08-23T22:12:01Z

Codecov Report

All modified and coverable lines are covered by tests ✅

Please upload report for BASE (v0.8.0-rc@d515d6c). Learn more about missing BASE report.

Additional details and impacted files

@@             Coverage Diff              @@
##             v0.8.0-rc    #1047   +/-   ##
============================================
  Coverage             ?   97.89%           
============================================
  Files                ?       75           
  Lines                ?    10737           
  Branches             ?     1239           
============================================
  Hits                 ?    10511           
  Misses               ?      178           
  Partials             ?       48

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚨 Try these New Features:

Flaky Tests Detection - Detect and resolve failed and flaky tests

paul0403 · 2024-08-26T15:58:29Z

This bug came from mismatched indices.

Consider the following IR we send to the runtime capi (compiled from the added test example):

define internal { ptr, ptr, i64, [1 x i64], [1 x i64] } @f(ptr %0, ptr %1, i64 %2, i64 %3, i64 %4) {
  %6 = insertvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } undef, ptr %0, 0
  %7 = insertvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %6, ptr %1, 1
  %8 = insertvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %7, i64 %2, 2
  %9 = insertvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %8, i64 %3, 3, 0
  %10 = insertvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %9, i64 %4, 4, 0
  call void @__catalyst__rt__device_init(ptr @"/home/paul.wang/catalyst_new/catalyst/frontend/catalyst/utils/../../../runtime/build/lib/librtd_lightning.so", ptr @LightningSimulator, ptr @"{'shots': 0, 'mcmc': False, 'num_burnin': 0, 'kernel_name': None}")
  %11 = call ptr @__catalyst__rt__qubit_allocate_array(i64 2)
  %12 = call ptr @__catalyst__rt__array_get_element_ptr_1d(ptr %11, i64 0)
  %13 = load ptr, ptr %12, align 8
  %14 = call ptr @__catalyst__rt__array_get_element_ptr_1d(ptr %11, i64 1)
  %15 = load ptr, ptr %14, align 8
  %16 = alloca { ptr, ptr, i64, [1 x i64], [1 x i64] }, i64 1, align 8
  store { ptr, ptr, i64, [1 x i64], [1 x i64] } %10, ptr %16, align 8
  call void (ptr, i64, ...) @__catalyst__qis__SetState(ptr %16, i64 2, ptr %13, ptr %15)
  %17 = call ptr @_mlir_memref_to_llvm_alloc(i64 ptrtoint (ptr getelementptr (double, ptr null, i32 4) to i64))
  %18 = insertvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } undef, ptr %17, 0
  %19 = insertvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %18, ptr %17, 1
  %20 = insertvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %19, i64 0, 2
  %21 = insertvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %20, i64 4, 3, 0
  %22 = insertvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %21, i64 1, 4, 0
  %23 = alloca { ptr, ptr, i64, [1 x i64], [1 x i64] }, i64 1, align 8
  store { ptr, ptr, i64, [1 x i64], [1 x i64] } %22, ptr %23, align 8
  call void (ptr, i64, ...) @__catalyst__qis__Probs(ptr %23, i64 2, ptr %13, ptr %15)
  call void @__catalyst__rt__qubit_release_array(ptr %11)
  call void @__catalyst__rt__device_release()
  ret { ptr, ptr, i64, [1 x i64], [1 x i64] } %22
}

define internal { ptr, ptr, i64, [1 x i64], [1 x i64] } @g(ptr %0, ptr %1, i64 %2, i64 %3, i64 %4) {
  %6 = insertvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } undef, ptr %0, 0
  %7 = insertvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %6, ptr %1, 1
  %8 = insertvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %7, i64 %2, 2
  %9 = insertvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %8, i64 %3, 3, 0
  %10 = insertvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %9, i64 %4, 4, 0
  call void @__catalyst__rt__device_init(ptr @"/home/paul.wang/catalyst_new/catalyst/frontend/catalyst/utils/../../../runtime/build/lib/librtd_lightning.so", ptr @LightningSimulator, ptr @"{'shots': 0, 'mcmc': False, 'num_burnin': 0, 'kernel_name': None}")
  %11 = call ptr @__catalyst__rt__qubit_allocate_array(i64 2)
  %12 = call ptr @__catalyst__rt__array_get_element_ptr_1d(ptr %11, i64 0)
  %13 = load ptr, ptr %12, align 8
  %14 = call ptr @__catalyst__rt__array_get_element_ptr_1d(ptr %11, i64 1)
  %15 = load ptr, ptr %14, align 8
  %16 = alloca { ptr, ptr, i64, [1 x i64], [1 x i64] }, i64 1, align 8
  store { ptr, ptr, i64, [1 x i64], [1 x i64] } %10, ptr %16, align 8
  call void (ptr, i64, ...) @__catalyst__qis__SetState(ptr %16, i64 2, ptr %13, ptr %15)
  %17 = call ptr @_mlir_memref_to_llvm_alloc(i64 ptrtoint (ptr getelementptr (double, ptr null, i32 4) to i64))
  %18 = insertvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } undef, ptr %17, 0
  %19 = insertvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %18, ptr %17, 1
  %20 = insertvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %19, i64 0, 2
  %21 = insertvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %20, i64 4, 3, 0
  %22 = insertvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %21, i64 1, 4, 0
  %23 = alloca { ptr, ptr, i64, [1 x i64], [1 x i64] }, i64 1, align 8
  store { ptr, ptr, i64, [1 x i64], [1 x i64] } %22, ptr %23, align 8
  call void (ptr, i64, ...) @__catalyst__qis__Probs(ptr %23, i64 2, ptr %13, ptr %15)
  call void @__catalyst__rt__qubit_release_array(ptr %11)
  call void @__catalyst__rt__device_release()
  ret { ptr, ptr, i64, [1 x i64], [1 x i64] } %22
}

The first call, __catalyst__rt__qubit_allocate_array, calls the AllocateQubits of the quantum device objects and (more or less) returns whatever AllocateQubits returns (https://github.com/PennyLaneAI/catalyst/blob/main/runtime/lib/capi/RuntimeCAPI.cpp#L325). The quantum device's AllocateQubits calls the device's QubitManager's AllocateRange (https://github.com/PennyLaneAI/catalyst/blob/main/runtime/lib/backend/lightning/lightning_dynamic/LightningSimulator.cpp#L39) and returns its results.

The AllocateRange will return [0,1] for the first call by f and [2,3] for the second call by g (which gets stored into the SSA value %11 in the above example), because it collects the QubitManager's next_idx, which does not get reset to zero when the manager is released (https://github.com/PennyLaneAI/catalyst/blob/main/runtime/lib/backend/common/QubitManager.hpp#L129).

Then, when __catalyst__qis__SetState is called with arguments %13, %15 ([0,1] for the first call, [2,3] for the second call), it passes these values down to the device's SetState (https://github.com/PennyLaneAI/catalyst/blob/main/runtime/lib/capi/RuntimeCAPI.cpp#L467). The device then uses these indices to try to set the statevector (https://github.com/PennyLaneAI/catalyst/blob/main/runtime/lib/backend/lightning/lightning_dynamic/LightningSimulator.cpp#L112), leading to "vector-index-out-of-range" segfaults, specifically at https://github.com/PennyLaneAI/pennylane-lightning/blob/97554407e114db41690b665c01f342c8ec27b98a/pennylane_lightning/core/src/simulators/lightning_qubit/StateVectorLQubit.hpp#L809

…tateprep

…tate

paul0403 · 2024-08-26T18:01:04Z

The root cause is because SetState (and BasisState) did not query the qubit_manager map values and just used the keys.

…tateprep

erick-xanadu

This makes sense! Thanks! I must have forgotten to add the getDeviceWires function call. All good on my part.

paul0403 · 2024-08-26T18:38:32Z

I guess we need to update the kokkos ones in lightning as well? https://github.com/PennyLaneAI/pennylane-lightning/blob/4c6e4bb750b7d976fe661677895692ef6c129723/pennylane_lightning/core/src/simulators/lightning_kokkos/catalyst/LightningKokkosSimulator.cpp#L140

Interestingly, the failure in #1044 does not happen on kokkos. The commit we track from lightning apparently doesn't even have the catalyst setstate and setbasisstate APIs

~~Since CI tracks lastest/stable/rc lightning (~~

catalyst/.github/workflows/check-pl-compat.yaml

Line 114 in 2f1398e

LIGHTNING_GIT_TAG_VALUE=master \

) ~~and CI passes, perhaps we don't need changes in lightning?~~
nvm, CI tracks that specific commit as well https://github.com/PennyLaneAI/catalyst/actions/runs/10564765994/job/29267907745?pr=1047#step:4:52

dime10 · 2024-08-27T15:00:16Z

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

@@ -108,14 +108,16 @@ void LightningSimulator::SetState(DataView<std::complex<double>, 1> &data,
                                  std::vector<QubitIdType> &wires)
 {
    std::vector<std::complex<double>> data_vector(data.begin(), data.end());
-    std::vector<std::size_t> wires_size_t(wires.begin(), wires.end());
+    auto &&dev_wires = getDeviceWires(wires);
+    std::vector<std::size_t> wires_size_t(dev_wires.begin(), dev_wires.end());


I don't think we need to repackage the wires into an additional vector

This was suggested on the lightning side too (PennyLaneAI/pennylane-lightning#869), I will change this

…kos (#869) **Context:** Catalyst recently fixed a bug where recycling a device when a previous execution involves stateprep causes a crash: PennyLaneAI/catalyst#1047 The fix is made to `lightning.qubit`, and thus we need to fix it for `lightning.kokkos` here as well. **Description of the Change:** In `LightningKokkosSimulator`, `SetState` and `SetBasisState` now correctly query the DevQubits from the SimQubits in its qubit_manger's map. **Benefits:** We can now have multiple qnode functions involving stateprep in a workflow **Possible Drawbacks:** **Related GitHub Issues:** PennyLaneAI/catalyst#1044 --------- Co-authored-by: ringo-but-quantum <github-ringo-but-quantum@xanadu.ai> Co-authored-by: Vincent Michaud-Rioux <vincentm@nanoacademic.com> Co-authored-by: Ali Asadi <10773383+maliasadi@users.noreply.github.com>

…r.cpp/`SetState` (and `SetBasisState`) (#1061) **Context:** Eliminating an unnecessary in-between variable in runtime/LightningSimulator.cpp/setstate and setbasisstate, when getting dev_wires from sim_wires in the qubit manager's map. See PennyLaneAI/pennylane-lightning#869 (comment), and #1047 (comment) **Benefits:** One less object initialization

…kos (#869) **Context:** Catalyst recently fixed a bug where recycling a device when a previous execution involves stateprep causes a crash: PennyLaneAI/catalyst#1047 The fix is made to `lightning.qubit`, and thus we need to fix it for `lightning.kokkos` here as well. **Description of the Change:** In `LightningKokkosSimulator`, `SetState` and `SetBasisState` now correctly query the DevQubits from the SimQubits in its qubit_manger's map. **Benefits:** We can now have multiple qnode functions involving stateprep in a workflow **Possible Drawbacks:** **Related GitHub Issues:** PennyLaneAI/catalyst#1044 --------- Co-authored-by: ringo-but-quantum <github-ringo-but-quantum@xanadu.ai> Co-authored-by: Vincent Michaud-Rioux <vincentm@nanoacademic.com> Co-authored-by: Ali Asadi <10773383+maliasadi@users.noreply.github.com>

…okkos (#878) **Context:** Catalyst recently fixed a bug where recycling a device when a previous execution involves stateprep causes a crash: PennyLaneAI/catalyst#1047 The fix is made to `lightning.qubit`, and thus we need to fix it for `lightning.kokkos` here as well. **Description of the Change:** In `LightningKokkosSimulator`, `SetState` and `SetBasisState` now correctly query the DevQubits from the SimQubits in its qubit_manger's map. **Benefits:** We can now have multiple qnode functions involving stateprep in a workflow **Possible Drawbacks:** **Related GitHub Issues:** PennyLaneAI/catalyst#1044 --------- Co-authored-by: paul0403 <79805239+paul0403@users.noreply.github.com> Co-authored-by: ringo-but-quantum <github-ringo-but-quantum@xanadu.ai> Co-authored-by: Ali Asadi <10773383+maliasadi@users.noreply.github.com>

Do not recycle device whatsoever solves the issue; we should record w…

bf027f9

…hether there is stateprep in a previous device to judge whether the device should be recycled.

paul0403 requested review from josh146, dime10 and erick-xanadu August 23, 2024 14:55

add TODO to SetState()

7c2bd25

paul0403 changed the title ~~Do not recycle device when a previous execution involves stateprep~~ Reset wire index when releasing devices Aug 23, 2024

resetting device wire index to zero upon device release

30a4cb9

frontend tests

bca1d0b

paul0403 commented Aug 23, 2024

View reviewed changes

runtime/lib/backend/common/QubitManager.hpp Outdated Show resolved Hide resolved

paul0403 added 2 commits August 26, 2024 09:39

Merge remote-tracking branch 'origin/main' into recycle_device_stateprep

b32cd7f

changelog

d911f0f

paul0403 added this to the v0.8.0 milestone Aug 26, 2024

dime10 added bug Something isn't working runtime Pull requests that update the runtime labels Aug 26, 2024

rauletorresc changed the base branch from main to v0.8.0-rc August 26, 2024 16:00

paul0403 added 2 commits August 26, 2024 12:25

Merge remote-tracking branch 'origin/v0.8.0-rc' into recycle_device_s…

68c928b

…tateprep

querying DevQubitIds from qubit_manager's map for setstate and basiss…

e1eec99

…tate

Merge remote-tracking branch 'origin/v0.8.0-rc' into recycle_device_s…

806bc25

…tateprep

paul0403 changed the title ~~Reset wire index when releasing devices~~ Correctly querying wire IDs for SetState and BasisState in runtime Aug 26, 2024

paul0403 changed the title ~~Correctly querying wire IDs for SetState and BasisState in runtime~~ Correctly querying wire IDs for SetState and SetBasisState in runtime Aug 26, 2024

erick-xanadu approved these changes Aug 26, 2024

View reviewed changes

paul0403 merged commit 68f579e into v0.8.0-rc Aug 26, 2024
40 checks passed

paul0403 deleted the recycle_device_stateprep branch August 26, 2024 19:09

paul0403 mentioned this pull request Aug 26, 2024

Correctly querying wire IDs for SetState and SetBasisState in kokkos PennyLaneAI/pennylane-lightning#869

Merged

5 tasks

dime10 reviewed Aug 27, 2024

View reviewed changes

paul0403 mentioned this pull request Aug 27, 2024

Remove unnecessary in-between variables for runtime/LightningSimulator.cpp/SetState (and SetBasisState) #1061

Merged

vincentmr mentioned this pull request Aug 28, 2024

Correctly querying wire IDs for SetState and SetBasisState in L-Kokkos PennyLaneAI/pennylane-lightning#878

Merged

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Correctly querying wire IDs for `SetState` and `SetBasisState` in runtime #1047

Correctly querying wire IDs for `SetState` and `SetBasisState` in runtime #1047

paul0403 commented Aug 23, 2024 •

edited

Loading

dime10 commented Aug 23, 2024

paul0403 commented Aug 23, 2024

paul0403 commented Aug 23, 2024 •

edited

Loading

paul0403 commented Aug 23, 2024

codecov bot commented Aug 23, 2024 •

edited

Loading

paul0403 commented Aug 26, 2024

paul0403 commented Aug 26, 2024

erick-xanadu left a comment

paul0403 commented Aug 26, 2024 •

edited

Loading

dime10 Aug 27, 2024

paul0403 Aug 27, 2024

Correctly querying wire IDs for SetState and SetBasisState in runtime #1047

Correctly querying wire IDs for SetState and SetBasisState in runtime #1047

Conversation

paul0403 commented Aug 23, 2024 • edited Loading

dime10 commented Aug 23, 2024

paul0403 commented Aug 23, 2024

paul0403 commented Aug 23, 2024 • edited Loading

paul0403 commented Aug 23, 2024

codecov bot commented Aug 23, 2024 • edited Loading

Codecov Report

paul0403 commented Aug 26, 2024

paul0403 commented Aug 26, 2024

erick-xanadu left a comment

Choose a reason for hiding this comment

paul0403 commented Aug 26, 2024 • edited Loading

dime10 Aug 27, 2024

Choose a reason for hiding this comment

paul0403 Aug 27, 2024

Choose a reason for hiding this comment

Correctly querying wire IDs for `SetState` and `SetBasisState` in runtime #1047

Correctly querying wire IDs for `SetState` and `SetBasisState` in runtime #1047

paul0403 commented Aug 23, 2024 •

edited

Loading

paul0403 commented Aug 23, 2024 •

edited

Loading

codecov bot commented Aug 23, 2024 •

edited

Loading

paul0403 commented Aug 26, 2024 •

edited

Loading