qjit function execution occurs asynchronously

Context

Thread-Level Speculation is a technique that has been used in various research to speed up general purpose programs by speculatively executing code downstream of a function call. The idea here is to do this in a similar manner to JAX, see Asynchronous Dispatch in the JAX docs.

JAX does not wait for the operation to complete before returning control to the Python program. Instead, JAX returns a DeviceArray value, which is a future, i.e., a value that will be produced in the future on an accelerator device but isn’t necessarily available immediately. Only when the value of the DeviceArray is queried is a blocking call generated.

Consider the following code snippet. Here, x , a device array returned as the result of evaluating f is a future DeviceArray, and blocking only occurs when a user requests the value of x in Python.

>>> f = qjit(hybrid_func)
>>> x = f(0.54)

Questions:

The assumption here is that this will lead to speedups in the following situation (this assumption needs to be validated, but should be apparent in an interpreted language):

>>> f = qjit(hybrid_func)
>>> g = qjit(hybrid_func2)
>>> x = f(0.54)
>>> y = g(x)

That is, since x is evaluated asynchronously, Python is not blocked awaiting the result of f and can simply invoke g directly.

Requirements:

1. Code downstream of a qjit'ted function is executed in parallel with the compiled function.
2. Parallel evaluation halts once an instruction is reached that depends upon the result of the qjit'ted function.
3. The downstream code is only executed once.

Installation Help

Refer to the Catalyst installation guide for how to install a source build of the project.

Start by reviewing the qjit call flow and the linked JAX Asynchronous Dispatch documentation to understand the proposed future-value behavior. Validate the three stated requirements: downstream execution runs in parallel, stops at dependencies, and executes only once; use the Catalyst source-build installation guide to prepare the environment.