catalyst
catalyst copied to clipboard
PennyLaneAI
Make async qnodes specific to qnodes, not the whole JIT context.
Instead of the async flag specifying async for all qnodes in the JIT context, we could have an async flag for specific qnodes if we move the flag to the qnode. E.g.,
Change:
@qjit(async_qnodes=True)
def jit_context():
...
To:
@qnode(async=True)
def circuit():
...
PennyLane might not want this change unless they can support asynchronous execution as well.
Another (minor) issue in your example is that async is a reserved keyword, which is why we went for async_qnodes.
Another (minor) issue in your example is that async is a reserved keyword, which is why we went for async_qnodes.
The keyword is trivial.
PennyLane might not want this change unless they can support asynchronous execution as well.
PennyLane can just ignore this keyword. Simliar to inline in C. It could be just a compiler suggestion.
Alternatively, we can come up with another UI/UX to specify compiler options like:
@opt(asyncExec=True)
@qml.qnode(...)
def circuit():
...
Potentially this is a wider conversation, because I think this goes beyond just async. Namely: how can we allow users to specify local compiler options (async, error mitigation, etc.) within the @qjit?
For example, we could go with a contextdecorator-style approach, allowing users to provide compiler directives by decorating whole functions, or very locally via with statements:
@qjit
def f(x):
@catalyst.pragma(mitigate="zne")
@qnode(dev)
def circuit(...):
...
with catalyst.pragma(async_exec=True):
y = circuit(x)
z = circuit(x ** 2)
@josh146 technically all of these are just transforms, they just happen in MLIR as opposed to python. We could just use the same convention (a decorator per transform) and the transform just sets a flag in the tape and returns the same result.
I agree that we could also investigate a contextdecorator style approach for a more fine-grained approach.
I like the context-style approach because the experience with Catalyst has taught me that needing functions for everything can be annoying.
For async option specifically, I would first ask though is this even useful? In JAX for example everything is run asynchronously and the user doesn't even know or need to know.
I like the context-style approach because the experience with Catalyst has taught me that needing functions for everything can be annoying.
I have the same experience 😆 the thing I like about context-decorators is you can satisfy both groups of users with a singular piece of logic/code.
For async option specifically, I would first ask though is this even useful? In JAX for example everything is run asynchronously and the user doesn't even know or need to know.
async is not always a guarantee that things will run fast. We can be smart and have heuristics for when we believe async can be beneficial and automatically apply async to those qnodes automatically. However, a user (including us) might be smarter and might want to enable / disable async for specific qnodes.