Recycling the same `device.so` by different `quantum.set_state` causes a crash #1044

paul0403 · 2024-08-22T15:34:15Z

The following mlir has multiple functions using set_state on the same device shared object file and it crashes:

module @circuit_pytest {
  func.func public @jit_circuit_pytest(%arg0: tensor<f64>, %arg1: tensor<3xf64>, %arg2: tensor<f64>) -> tensor<f64> attributes {llvm.emit_c_interface} {
    %cst = arith.constant dense<[(1.000000e+00,0.000000e+00), (0.000000e+00,0.000000e+00), (0.000000e+00,0.000000e+00), (0.000000e+00,0.000000e+00)]> : tensor<4xcomplex<f64>>
    %0 = call @circuit_pytest(%cst, %arg0, %arg1, %arg2) : (tensor<4xcomplex<f64>>, tensor<f64>, tensor<3xf64>, tensor<f64>) -> tensor<f64>
    return %0 : tensor<f64>
  }

  func.func @circuit_pytest_tape_2(%arg0: tensor<4xcomplex<f64>>, %arg1: tensor<3xf64>) -> tensor<f64> {
    quantum.device["/home/paul.wang/catalyst_new/catalyst/frontend/catalyst/utils/../../../runtime/build/lib/librtd_lightning.so", "LightningSimulator", "{'shots': 0, 'mcmc': False, 'num_burnin': 0, 'kernel_name': None}"]
    %0 = quantum.alloc( 2) : !quantum.reg
    %1 = quantum.extract %0[ 0] : !quantum.reg -> !quantum.bit
    %2 = quantum.extract %0[ 1] : !quantum.reg -> !quantum.bit
    %3:2 = quantum.set_state(%arg0) %1, %2 : (tensor<4xcomplex<f64>>, !quantum.bit, !quantum.bit) -> (!quantum.bit, !quantum.bit)
    %extracted_slice = tensor.extract_slice %arg1[2] [1] [1] : tensor<3xf64> to tensor<1xf64>
    %collapsed = tensor.collapse_shape %extracted_slice [] : tensor<1xf64> into tensor<f64>
    %extracted = tensor.extract %collapsed[] : tensor<f64>
    %out_qubits = quantum.custom "RX"(%extracted) %3#1 : !quantum.bit
    %4 = quantum.namedobs %3#0[ PauliZ] : !quantum.obs
    %5 = quantum.expval %4 : f64
    %from_elements = tensor.from_elements %5 : tensor<f64>
    %6 = quantum.insert %0[ 0], %3#0 : !quantum.reg, !quantum.bit
    %7 = quantum.insert %6[ 1], %out_qubits : !quantum.reg, !quantum.bit
    quantum.dealloc %7 : !quantum.reg
    quantum.device_release
    return %from_elements : tensor<f64>
  }

  func.func @circuit_pytest_tape_1(%arg0: tensor<4xcomplex<f64>>, %arg1: tensor<3xf64>) -> tensor<f64> {
    quantum.device["/home/paul.wang/catalyst_new/catalyst/frontend/catalyst/utils/../../../runtime/build/lib/librtd_lightning.so", "LightningSimulator", "{'shots': 0, 'mcmc': False, 'num_burnin': 0, 'kernel_name': None}"]
    %0 = quantum.alloc( 2) : !quantum.reg
    %1 = quantum.extract %0[ 0] : !quantum.reg -> !quantum.bit
    %2 = quantum.extract %0[ 1] : !quantum.reg -> !quantum.bit
    %3:2 = quantum.set_state(%arg0) %1, %2 : (tensor<4xcomplex<f64>>, !quantum.bit, !quantum.bit) -> (!quantum.bit, !quantum.bit)
    %4 = quantum.namedobs %3#0[ PauliZ] : !quantum.obs
    %5 = quantum.expval %4 : f64
    %from_elements = tensor.from_elements %5 : tensor<f64>
    %6 = quantum.insert %0[ 0], %3#0 : !quantum.reg, !quantum.bit
    %extracted_slice = tensor.extract_slice %arg1[1] [1] [1] : tensor<3xf64> to tensor<1xf64>
    %collapsed = tensor.collapse_shape %extracted_slice [] : tensor<1xf64> into tensor<f64>
    %extracted = tensor.extract %collapsed[] : tensor<f64>
    %out_qubits = quantum.custom "RX"(%extracted) %3#1 : !quantum.bit
    %7 = quantum.insert %6[ 1], %out_qubits : !quantum.reg, !quantum.bit
    quantum.dealloc %7 : !quantum.reg
    quantum.device_release
    return %from_elements : tensor<f64>
  }

  func.func @circuit_pytest_tape_0(%arg0: tensor<4xcomplex<f64>>, %arg1: tensor<3xf64>) -> tensor<f64> {
    quantum.device["/home/paul.wang/catalyst_new/catalyst/frontend/catalyst/utils/../../../runtime/build/lib/librtd_lightning.so", "LightningSimulator", "{'shots': 0, 'mcmc': False, 'num_burnin': 0, 'kernel_name': None}"]
    %0 = quantum.alloc( 2) : !quantum.reg
    %1 = quantum.extract %0[ 0] : !quantum.reg -> !quantum.bit
    %2 = quantum.extract %0[ 1] : !quantum.reg -> !quantum.bit
    %3:2 = quantum.set_state(%arg0) %1, %2 : (tensor<4xcomplex<f64>>, !quantum.bit, !quantum.bit) -> (!quantum.bit, !quantum.bit)
    %4 = quantum.namedobs %3#0[ PauliZ] : !quantum.obs
    %5 = quantum.expval %4 : f64
    %from_elements = tensor.from_elements %5 : tensor<f64>
    %6 = quantum.insert %0[ 0], %3#0 : !quantum.reg, !quantum.bit
    %extracted_slice = tensor.extract_slice %arg1[0] [1] [1] : tensor<3xf64> to tensor<1xf64>
    %collapsed = tensor.collapse_shape %extracted_slice [] : tensor<1xf64> into tensor<f64>
    %extracted = tensor.extract %collapsed[] : tensor<f64>
    %out_qubits = quantum.custom "RX"(%extracted) %3#1 : !quantum.bit
    %7 = quantum.insert %6[ 1], %out_qubits : !quantum.reg, !quantum.bit
    quantum.dealloc %7 : !quantum.reg
    quantum.device_release
    return %from_elements : tensor<f64>
  }

  func.func private @circuit_pytest(%arg0: tensor<4xcomplex<f64>>, %arg1: tensor<f64>, %arg2: tensor<3xf64>, %arg3: tensor<f64>) -> tensor<f64> attributes {diff_method = "parameter-shift", llvm.linkage = #llvm.linkage<internal>, qnode} {
    %0 = call @circuit_pytest_tape_0(%arg0, %arg2) : (tensor<4xcomplex<f64>>, tensor<3xf64>) -> tensor<f64>
    %1 = call @circuit_pytest_tape_1(%arg0, %arg2) : (tensor<4xcomplex<f64>>, tensor<3xf64>) -> tensor<f64>
    %2 = call @circuit_pytest_tape_2(%arg0, %arg2) : (tensor<4xcomplex<f64>>, tensor<3xf64>) -> tensor<f64>
    return %0 : tensor<f64>
  }

  func.func @setup() {
    quantum.init
    return
  }
  func.func @teardown() {
    quantum.finalize
    return
  }
}

malloc(): unaligned tcache chunk detected

However, if we do one of the two following, then there won't be a crash:

create copies of the device by cp librtd_lightning.so librtd_lightning_1.so && cp librtd_lightning.so librtd_lightning_2.so, and use these copies in quantum.device
remove all the quantum.set_state, i.e.

module @circuit_pytest {
  func.func public @jit_circuit_pytest(%arg0: tensor<f64>, %arg1: tensor<3xf64>, %arg2: tensor<f64>) -> tensor<f64> attributes {llvm.emit_c_interface} {
    %cst = arith.constant dense<[(1.000000e+00,0.000000e+00), (0.000000e+00,0.000000e+00), (0.000000e+00,0.000000e+00), (0.000000e+00,0.000000e+00)]> : tensor<4xcomplex<f64>>
    %0 = call @circuit_pytest(%cst, %arg0, %arg1, %arg2) : (tensor<4xcomplex<f64>>, tensor<f64>, tensor<3xf64>, tensor<f64>) -> tensor<f64>
    return %0 : tensor<f64>
  }

  func.func @circuit_pytest_tape_2(%arg0: tensor<4xcomplex<f64>>, %arg1: tensor<3xf64>) -> tensor<f64> {
    quantum.device["/home/paul.wang/catalyst_new/catalyst/frontend/catalyst/utils/../../../runtime/build/lib/librtd_lightning.so", "LightningSimulator", "{'shots': 0, 'mcmc': False, 'num_burnin': 0, 'kernel_name': None}"]
    %0 = quantum.alloc( 2) : !quantum.reg
    %1 = quantum.extract %0[ 0] : !quantum.reg -> !quantum.bit
    %2 = quantum.extract %0[ 1] : !quantum.reg -> !quantum.bit
    //%3:2 = quantum.set_state(%arg0) %1, %2 : (tensor<4xcomplex<f64>>, !quantum.bit, !quantum.bit) -> (!quantum.bit, !quantum.bit)
    %extracted_slice = tensor.extract_slice %arg1[2] [1] [1] : tensor<3xf64> to tensor<1xf64>
    %collapsed = tensor.collapse_shape %extracted_slice [] : tensor<1xf64> into tensor<f64>
    %extracted = tensor.extract %collapsed[] : tensor<f64>
    //%out_qubits = quantum.custom "RX"(%extracted) %3#1 : !quantum.bit
    %out_qubits = quantum.custom "RX"(%extracted) %2 : !quantum.bit
    //%4 = quantum.namedobs %3#0[ PauliZ] : !quantum.obs
    %4 = quantum.namedobs %1[ PauliZ] : !quantum.obs
    %5 = quantum.expval %4 : f64
    %from_elements = tensor.from_elements %5 : tensor<f64>
    //%6 = quantum.insert %0[ 0], %3#0 : !quantum.reg, !quantum.bit
    %6 = quantum.insert %0[ 0], %1 : !quantum.reg, !quantum.bit
    %7 = quantum.insert %6[ 1], %out_qubits : !quantum.reg, !quantum.bit
    quantum.dealloc %7 : !quantum.reg
    quantum.device_release
    return %from_elements : tensor<f64>
  }

  func.func @circuit_pytest_tape_1(%arg0: tensor<4xcomplex<f64>>, %arg1: tensor<3xf64>) -> tensor<f64> {
    quantum.device["/home/paul.wang/catalyst_new/catalyst/frontend/catalyst/utils/../../../runtime/build/lib/librtd_lightning.so", "LightningSimulator", "{'shots': 0, 'mcmc': False, 'num_burnin': 0, 'kernel_name': None}"]
    %0 = quantum.alloc( 2) : !quantum.reg
    %1 = quantum.extract %0[ 0] : !quantum.reg -> !quantum.bit
    %2 = quantum.extract %0[ 1] : !quantum.reg -> !quantum.bit
    //%3:2 = quantum.set_state(%arg0) %1, %2 : (tensor<4xcomplex<f64>>, !quantum.bit, !quantum.bit) -> (!quantum.bit, !quantum.bit)
    //%4 = quantum.namedobs %3#0[ PauliZ] : !quantum.obs
    %4 = quantum.namedobs %1[ PauliZ] : !quantum.obs
    %5 = quantum.expval %4 : f64
    %from_elements = tensor.from_elements %5 : tensor<f64>
    //%6 = quantum.insert %0[ 0], %3#0 : !quantum.reg, !quantum.bit
    %6 = quantum.insert %0[ 0], %1 : !quantum.reg, !quantum.bit
    %extracted_slice = tensor.extract_slice %arg1[1] [1] [1] : tensor<3xf64> to tensor<1xf64>
    %collapsed = tensor.collapse_shape %extracted_slice [] : tensor<1xf64> into tensor<f64>
    %extracted = tensor.extract %collapsed[] : tensor<f64>
    //%out_qubits = quantum.custom "RX"(%extracted) %3#1 : !quantum.bit
    %out_qubits = quantum.custom "RX"(%extracted) %2 : !quantum.bit
    %7 = quantum.insert %6[ 1], %out_qubits : !quantum.reg, !quantum.bit
    quantum.dealloc %7 : !quantum.reg
    quantum.device_release
    return %from_elements : tensor<f64>
  }

  func.func @circuit_pytest_tape_0(%arg0: tensor<4xcomplex<f64>>, %arg1: tensor<3xf64>) -> tensor<f64> {
    quantum.device["/home/paul.wang/catalyst_new/catalyst/frontend/catalyst/utils/../../../runtime/build/lib/librtd_lightning.so", "LightningSimulator", "{'shots': 0, 'mcmc': False, 'num_burnin': 0, 'kernel_name': None}"]
    %0 = quantum.alloc( 2) : !quantum.reg
    %1 = quantum.extract %0[ 0] : !quantum.reg -> !quantum.bit
    %2 = quantum.extract %0[ 1] : !quantum.reg -> !quantum.bit
    //%3:2 = quantum.set_state(%arg0) %1, %2 : (tensor<4xcomplex<f64>>, !quantum.bit, !quantum.bit) -> (!quantum.bit, !quantum.bit)
    //%4 = quantum.namedobs %3#0[ PauliZ] : !quantum.obs
    %4 = quantum.namedobs %1[ PauliZ] : !quantum.obs
    %5 = quantum.expval %4 : f64
    %from_elements = tensor.from_elements %5 : tensor<f64>
    //%6 = quantum.insert %0[ 0], %3#0 : !quantum.reg, !quantum.bit
    %6 = quantum.insert %0[ 0], %1 : !quantum.reg, !quantum.bit
    %extracted_slice = tensor.extract_slice %arg1[0] [1] [1] : tensor<3xf64> to tensor<1xf64>
    %collapsed = tensor.collapse_shape %extracted_slice [] : tensor<1xf64> into tensor<f64>
    %extracted = tensor.extract %collapsed[] : tensor<f64>
    //%out_qubits = quantum.custom "RX"(%extracted) %3#1 : !quantum.bit
    %out_qubits = quantum.custom "RX"(%extracted) %2 : !quantum.bit
    %7 = quantum.insert %6[ 1], %out_qubits : !quantum.reg, !quantum.bit
    quantum.dealloc %7 : !quantum.reg
    quantum.device_release
    return %from_elements : tensor<f64>
  }

  func.func private @circuit_pytest(%arg0: tensor<4xcomplex<f64>>, %arg1: tensor<f64>, %arg2: tensor<3xf64>, %arg3: tensor<f64>) -> tensor<f64> attributes {diff_method = "parameter-shift", llvm.linkage = #llvm.linkage<internal>, qnode} {
    %0 = call @circuit_pytest_tape_0(%arg0, %arg2) : (tensor<4xcomplex<f64>>, tensor<3xf64>) -> tensor<f64>
    %1 = call @circuit_pytest_tape_1(%arg0, %arg2) : (tensor<4xcomplex<f64>>, tensor<3xf64>) -> tensor<f64>
    %2 = call @circuit_pytest_tape_2(%arg0, %arg2) : (tensor<4xcomplex<f64>>, tensor<3xf64>) -> tensor<f64>
    return %0 : tensor<f64>
  }
  func.func @setup() {
    quantum.init
    return
  }
  func.func @teardown() {
    quantum.finalize
    return
  }
}

The text was updated successfully, but these errors were encountered:

paul0403 · 2024-08-23T13:43:20Z

A minimal frontend example to reproduce this error (as of main branch b405fbb):

dev = qml.device("lightning.qubit", wires=2)

@qml.qnode(dev)
def f():
    qml.StatePrep(
                np.array([complex(1, 0), complex(0, 0), complex(0, 0), complex(0, 0)]),
                wires=[0, 1],
            )
    return qml.probs()


@qml.qnode(dev)
def g():
    qml.StatePrep(
                np.array([complex(1, 0), complex(0, 0), complex(0, 0), complex(0, 0)]),
                wires=[0, 1],
            )
    return qml.probs()

@qjit
def main():
    return f(), g()

main()

malloc(): unaligned tcache chunk detected

Quick note: this is an mlir issue and not a jax issue so I know this wouldn't help, but I checked it anyway, and swapping out jax.numpy.array for np.array results in the same error.

… crash PennyLaneAI/catalyst#1044

…time (#1047) **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]

…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>

paul0403 added the bug Something isn't working label Aug 23, 2024

paul0403 mentioned this issue Aug 23, 2024

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

Merged

paul0403 added a commit to PennyLaneAI/pennylane-lightning that referenced this issue Aug 23, 2024

Recycling the same device.so by different SetStateVector causes a…

7536166

… crash PennyLaneAI/catalyst#1044

paul0403 mentioned this issue Aug 24, 2024

Split multiple tapes into multiple functions in mlir #1017

Merged

paul0403 mentioned this issue Aug 26, 2024

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

Merged

5 tasks

vincentmr mentioned this issue Aug 28, 2024

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

Merged

paul0403 closed this as completed Sep 4, 2024

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Recycling the same `device.so` by different `quantum.set_state` causes a crash #1044

Recycling the same `device.so` by different `quantum.set_state` causes a crash #1044

paul0403 commented Aug 22, 2024 •

edited

Loading

paul0403 commented Aug 23, 2024 •

edited

Loading

Recycling the same device.so by different quantum.set_state causes a crash #1044

Recycling the same device.so by different quantum.set_state causes a crash #1044

Comments

paul0403 commented Aug 22, 2024 • edited Loading

paul0403 commented Aug 23, 2024 • edited Loading

Recycling the same `device.so` by different `quantum.set_state` causes a crash #1044

Recycling the same `device.so` by different `quantum.set_state` causes a crash #1044

paul0403 commented Aug 22, 2024 •

edited

Loading

paul0403 commented Aug 23, 2024 •

edited

Loading