catalyst
catalyst copied to clipboard
Python only JIT-like proof of concept
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:
- Replace qnode calls with
pure_callbacks. - Inside the
pure_callbackwe have a template (at the moment it is hard coded) of a quantum function - We compile with Catalyst and link it again
- We run it.
- 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.
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.
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
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)