[qref 3.8] Convert subroutines in plxpr to qref catxpr

[paul0403]

Context:
Convert subroutines in plxpr to qref catxpr.

The use of QubitHandler is removed, since both plxpr and catalyst jaxpr is now reference semantics.

When calling a subroutine, there are some different cases. In python, a subroutine is always called with "wires", which can either be integers or dynamically allocated wires:

@qp.capture.subroutine
def subroutine(wires):
     qp.Toffoli(wires=[1, *wires])

@qp.qjit(capture=True, target="mlir")
@qp.qnode(qp.device("lightning.qubit", wires=3))
def circuit(i: int):
    with qp.allocate(2) as q:
        subroutine_with_allocation(wires=[q[0], 0])   # call with an allocated wire and an int

In plxpr, dynamic allocations return a specific AbstractQubit type instead of a generic integer type (PennyLaneAI/pennylane#9400). This means in the signature of the subroutine function (and the corresponding calls), the arguments can either be integers or an abstract qubit type.

When the arguments are integers, we must not implicitly convert the func args to qref.bit types, because the subroutine function can use them as actual integers as well. Therefore, on top of the integer, the global register must also be taken in.

When arguments are AbstractQubit types, the args can directly be qref.bit types. Therefore the above will compile to

    func.func public @circuit(%arg0: tensor<i64>){
      ....
      %0 = qref.alloc( 3) : !qref.reg<3>
      %1 = qref.alloc( 2) : !qref.reg<2>
      %2 = qref.get %1[ 0] : !qref.reg<2> -> !qref.bit
      call @subroutine_with_allocation(%0, %2, %c) : (!qref.reg<3>, !qref.bit, tensor<i64>) -> ()
      qref.dealloc %1 : !qref.reg<2>
      ....
    }
    func.func private @subroutine(%arg0: !qref.reg<3>, %arg1: !qref.bit, %arg2: tensor<i64>) {
      %0 = qref.get %arg0[ 1] : !qref.reg<3> -> !qref.bit
      %extracted = tensor.extract %arg2[] : tensor<i64>
      %1 = qref.get %arg0[%extracted] : !qref.reg<3>, i64 -> !qref.bit   // get global wire index 0
      qref.custom "Toffoli"() %0, %arg1, %1 : !qref.bit, !qref.bit, !qref.bit
      return
    }

In this scheme, inside the subroutine:

• aliasing between any global wires are handled automatically because any global wire indices come in as an integer arg + the global qreg arg at the beginning, and they are automatically dynamic
• there will be no alisaing between global wires and dynamically allocated wires because they belong in the two different type groups
• The only potential alias is between two qref.bit args. When both extract indices are dynamic, we directly assert that they are different inside the value semantics conversion pass. In cases of dynamic indices, we add an catalyst.assert op to check they are different in runtime.

[sc-115874]

[rniczh]

💯