CUDA Quantum Interpreter: Provide semantics for `for_p`.

Context

Catalyst has recently added support for executing quantum programs in NVIDIA's CUDA Quantum platform. For example, in the following code, we see two identical quantum programs. Both programs will execute the RX gate with parameter a, and will return the state of the system.

import pennylane as qml
from catalyst import qjit
from catalyst.cuda import qjit as cjit, SoftwareQQPP

@qjit
@qml.qnode(qml.device("lightning.qubit", wires=1))
def foo(a):
   qml.RX(a, wires=0)
   return qml.state()

@cjit
@qml.qnode(qml.device(SoftwareQQPP(wires=1)))
def bar(a):
   qml.RX(a, wires=0)
   return qml.state()

These equivalent quantum programs are running on different simulators. The first one, has been specified to run in the lightning.qubit simulator, while the second one has been specified to run in the qpp-cpu simulator

These equivalent quantum programs are written using the PennyLane's API. However, in order to execute these programs in the qpp-cpu simulator, we first need to translate them into NVIDIA's CUDA Quantum Python API to describe quantum programs. The program above written in NVIDIA's CUDA Quantum Python API could look like the following:
import cudaq

def bar(a):
  kernel = cudaq.make_kernel()
  qreg = kernel.qalloc(1)
  qubit0 = qreg[0]
  kernel.rx(a, qubit0)
  return cudaq.get_state(kernel)

Goal

Support for translating quantum programs written in PennyLane's API into NVIDIA's CUDA Quantum Python API is limited at the moment. In particular, we don't have support for translating PennyLane's for loops statements in CUDA Quantum. Here is how one would express conditional statements and for loops in PennyLane.

        @qjit()
        @qml.qnode(qml.device(backend, wires=6))
        def circuit(n: int):
            qml.Hadamard(wires=0)

            @for_loop(0, n - 1, 1)
            def loop_fn(i):
                qml.CNOT(wires=[i, i + 1])

            loop_fn()
            return qml.state()

CUDA Quantum's Python API also allows users to specify for loops.

Instead the above program should be translated to the following CUDA Quantum's Python API calls.

def circuit(n: int):

  kernel = cudaq.make_kernel()
  qreg = kernel.qalloc(7)
  qubit0 = qreg[0]
  kernel.h(qubit0)
  def loop(index):
    qubit_i = qreg[index]
    qubit_i_plus_1 = qreg[index + 1]
    kernel.cx(qubit_i, qubit_i_plus_1)
  kernel.for_loop(start=0, stop=n-1, function=loop)
  return cudaq.get_state(kernel)

Technical details

• PennyLane's API calls are converted to CUDA Quantum's Python API via a custom JAX interpreter. found in catalyst.cuda.catalyst_to_cuda_interpreter.py.
• You will need to implement the semantics for the Catalyst's JAX primitive for_p.
• Write a function that takes an InterprereterContext and a for_p equation and checks for the parameter to for_p.
• Construct a kernel.for_loop call that matches the semantics of for_p.
• You may want to construct a JAX primitive for for_loop similar to other CUDA Quantum's JAX primitives found in catalyst.cuda.catalyst.primitives along with convenience functions.

Start with catalyst.cuda.catalyst_to_cuda_interpreter.py and inspect how existing CUDA Quantum JAX primitives are defined in catalyst.cuda.catalyst.primitives. Trace the Catalyst for_p equation parameters and implement its interpreter semantics, including a matching kernel.for_loop call. Done means PennyLane for_loop programs translate into equivalent CUDA Quantum API calls.