Support for `IsingZZ` in Catalyst

[rmoyard]

Important Note

⚠️ This issue is part of an internal assignment and not meant for external contributors.

Context

We would like to add support for the PennyLane IsingZZ operation to our compilation architecture. The PennyLane-Lightning simulator supports this gate, and we would like to be able to support it without relying on the operation decomposition.

Goal

We would like to support the following example without having to rely on decomposition:

import pennylane as qml
from catalyst import *

dev = qml.device("lightning.qubit", wires=2, shots=100)

@qjit
@qml.qnode(dev)
def f(x: float):
    qml.Hadamard(wires=0)
    qml.IsingZZ(x, wires=[0, 1])
    return qml.probs()

Right now, indeed we see that by inspecting the captured program in the frontend: the decomposition occurs.

>>> f.jaxpr
{ lambda ; a:f64[]. let
    b:f64[4] = func[
      call_jaxpr={ lambda ; c:f64[]. let
          qdevice[
            rtd_kwargs={'shots': 100, 'mcmc': False, 'num_burnin': 0, 'kernel_name': None}
            rtd_lib=/home/romain/Catalyst/catalyst/frontend/catalyst/utils/../../../runtime/build/lib/librtd_lightning.so
            rtd_name=LightningSimulator
          ] 
          d:AbstractQreg() = qalloc 2
          e:AbstractQbit() = qextract d 0
          f:AbstractQbit() = qinst[
            ctrl_len=0
            op=Hadamard
            params_len=0
            qubits_len=1
          ] e
          g:AbstractQbit() = qextract d 1
          h:AbstractQbit() i:AbstractQbit() = qinst[
            ctrl_len=0
            op=CNOT
            params_len=0
            qubits_len=2
          ] f g
          j:AbstractQbit() = qinst[ctrl_len=0 op=RZ params_len=1 qubits_len=1] i
            c
          k:AbstractQbit() l:AbstractQbit() = qinst[
            ctrl_len=0
            op=CNOT
            params_len=0
            qubits_len=2
          ] h j
          m:AbstractObs(num_qubits=2,primitive=compbasis) = compbasis k l
          n:f64[4] = probs[shape=(4,)] m
          o:AbstractQreg() = qinsert d 0 k
          p:AbstractQreg() = qinsert o 1 l
          qdealloc p
        in (n,) }
      fn=<QNode: device='<lightning.qubit device (wires=2, shots=100) at 0x76eeee15ba00>', interface='auto', diff_method='best'>
    ] a
  in (b,) }

IsingZZ is not present, instead we see two CNOT and one RZ, which corresponds to the decomposition of the operation in the frontend.

Technical details

Catalyst handles unsupported gates in the Python layer (frontend). Catalyst decomposes any gates that are unsupported by the compilation pipeline, by the QIR interface between compiled program and runtime.

In order to support the new gates, they need to be added in several places:

• the list of supported gates in the qjit_device.py file under the RUNTIME_OPERATION list: https://github.com/PennyLaneAI/catalyst/pull/638/files#diff-49c4bb20bf2dcbd5fe316c548577fc168c4da1b075304b5bb953afcf8c4c7c25R37

• a QIR-style function declaration for each gate needs to be added to the C-API of the runtime: https://github.com/PennyLaneAI/catalyst/blob/cc580544ead54403ed054da75a6694ecdfc9a0b9/runtime/include/RuntimeCAPI.h#L51-L52

• an implementation of the C-API function, calling the QuantumDevice's NamedOperation function, similar to: https://github.com/PennyLaneAI/catalyst/blob/cc580544ead54403ed054da75a6694ecdfc9a0b9/runtime/lib/capi/RuntimeCAPI.cpp#L420

Requirements

A PR must satisfy the following:

• Add a lit test that shows that the decomposition does not occur: frontend/test/lit/test_multi_qubit_gates.py
• Add a unit test for NamedOperation("IsingZZ", ...) in runtime/tests/Test_LightningGateSet.cpp
• Add a unit test for __catalyst__qis__IsingZZ in /runtime/tests/Test_LightningCoreQIS.cpp
• Pass format checks
• Include a changelog entry
• Note that integration tests are already included and will fail if something went wrong: frontend/test/pytest/test_operations.py