catalyst icon indicating copy to clipboard operation
catalyst copied to clipboard

Python only JIT-like proof of concept

Open erick-xanadu opened this issue 4 months ago • 3 comments

Context: This is not a fully fledge runtime function specialization pass, but it is close. These changes show a proof of concept that would allow internal researchers to prototype function specialization passes easily from within Python.

Description of the Change: Leveraging callbacks and the xDSL program representation we do the following:

  1. Replace qnode calls with pure_callbacks.
  2. Inside the pure_callback we have a template (at the moment it is hard coded) of a quantum function
  3. We compile with Catalyst and link it again
  4. We run it.
  5. We return these values from the callback.

Benefits: A small proof of concept of (not yet) a runtime function specialization pass, but almost there.

Possible Drawbacks:

TODOs:

  • [ ] Obtain the program representation from module instead of using a hard coded string
  • [ ] Use the program representation above as a template
  • [ ] Embed the runtime values as constants in the template above
  • [ ] Automatically detect the return type signature
  • [ ] May not work well with multiple executions
  • [ ] May depend on implementations of dlopen (tested on ubuntu).

This is definitely just a proof of concept. A more full fledged implementation using proper tools would be necessary for production.

Related GitHub Issues:

  • Notes:

  • Uses PennyLane PR rebased with master.

erick-xanadu avatar Jul 23 '25 14:07 erick-xanadu

Looking at this code now, I'm not quite sure what the purpose of the xdsl transform is. If we just want to wrap the qnode call into a callback that compiles the given we can probably do that while skipping over the transform part. I'm going to try that and see if we can also get rid of the hardcoded ir.

dime10 avatar Aug 05 '25 18:08 dime10

This version is quite a bit simpler and also runs on the latest main branch of catalyst, no specific PL version needed (with some smaller fixes like automating paths): https://github.com/PennyLaneAI/catalyst/commit/bbcc154504a4c9d02918e6fff9ad941743d96360

dime10 avatar Aug 05 '25 19:08 dime10

A significant chunk of the todos were addressed here: https://github.com/PennyLaneAI/catalyst/commit/07d17d5163e737eedfcb692b0d39b6c6f3a8515a

In particular, the IR for the qnode is no longer hardcoded. (see commit for more details)

dime10 avatar Aug 05 '25 21:08 dime10