Mixed State GPU acceleration Pennylane

PennyLane Help
1

Hi,

I have been using the “default.mixed” backend to conduct error correction experiments. However, I have been looking to speed up my circuits using GPUs. To my knowledge, there is no support for mixed-state GPU acceleration using a Pennylane. I believe there is an IBM-qiskit backend that does, but this requires some work arounds to use for my work.

I checked out this post:

While it does say that there was an updated with regards to back propogation, I could not find any other information on implementation for anything other than back progogation(such as operations). If this has not been implemented yet, I would like to take a crack at implementing backend changes to support GPU acceleration.

With Thanks,

Connor

2

Hi @connor_gambla ,

As you noticed mixed-state GPU acceleration is not currently available in PennyLane. I’m unaware of whether or not there’s an IBM backend that supports this.

Regarding the post that you mention, backpropagation is indeed available for PennyLane’s default.mixed device.

Regarding your proposal for implementing backend changes to support GPU acceleration, I think this might be quite a big endeavour. If you want to try it, I think the best would be to consider this as an external library based on PennyLane, that basically acts as an optional plugin. This way you can work on it independently and you wouldn’t have to worry about merging it into PennyLane, which is unlikely to work out at the moment.

If you decide to add a demo in your repo showing how to use this feature, we could then possibly market it on social media and add it to our community demos page! What do you think?

4

Awesome! Looking forward to seeing your progress on this.

3

Currently working on it. I believe there might be some tensor support that was added to default.mixed_state, which would allow it to run using GPU acceleration with minimal changes. Hopefully I can come up with something, and would be glad to let you guys know if I make any progress on it.

Best

5

Hi Catalina,

So I did a bit of digging the documentation. Turns out the default.mixed supports jax in its backend, which means that it can support GPU acceleration! Though by no means as optimized as something like lightning, it still offers a significant speed up. We can utilize jax.jit to accelerate the circuit the same way a circuit with something like state vector backends. Below is the link to a repo that contains a brief demonstration. Let me know what you guys think about the repo:

Link to demonstration: GitHub - cwgambla/Accelerating-Mixed-State-Simulation-Pennylane-With-GPUs

As far backend changes to default.mixed, there seems to be little more that I could do that would further accelerate things since the backend is already jax compatible. Building out a mixed state equivalent for lightning seems like a great long term project, but does not seem like something I could do by myself nor in a reasonable amount of time. However, if you guys are looking for some additional assistance in building out that or any other feature, please feel to reach out.

Thank you for your assistance, and I hope this helps.

Best,

Connor

6

Wow this is amazing, thanks for sharing it @connor_gambla !

We’ll take a deeper look at your code but from what I can see it could make a good community demo.

For the README it might be useful to add just some versioning details for anyone wanting to use this in the future:

  • Python version
  • JAX version
  • PennyLane version
  • Qiskit and Qiskit-aer versions
  • What kind of GPU you used for the tests