feat: add cluster visualization for FuncOp and DeviceInitOp operations

[andrijapau]

This PR adds the ability to view multiple QNodes in a workflow with clusters and their devices as nodes within said clusters.

Some examples,

dev1 = qml.device("lightning.qubit", wires=5)
dev2 = qml.device("null.qubit", wires=5)

@qml.qnode(dev1)
def my_qnode2():
    qml.X(0)

@qml.qnode(dev2)
def my_qnode1():
    qml.H(0)

@qml.qjit(autograph=True, target="mlir")
def my_workflow():
    my_qnode1()
    my_qnode2()

@qml.qjit(autograph=True, target="mlir")
@qml.qnode(dev2)
def circuit():
    qml.H(0)

[sc-104730]
[sc-104731]

[mudit2812]

What if there is only 1 qnode but there are also classical instructions outside the qnode call? I.e.

@qml.qjit
def workflow(x):
    <classical operations>

    @qml.qnode(dev)
    def circuit(...):
        ...

    res = circuit(...)
    <classical stuff>
    return ...

[andrijapau]

Good question. For now we are not visualizing classical operations so only quantum operations within the circuit will render. The way the visualization is currently written, it will only show a single circuit cluster but perhaps it should also have the outer workflow cluster too ... 🤔

CC @AntonNI8 - what do you think?

[AntonNI8]

It's fine if we have an outer workflow cluster, but it's kept blank (the only thing inside it is the QNode). That should allow a smoother transition to multi-QNode workflows later on.

[andrijapau]

As discussed with @AntonNI8 here, we will have an outer bounding box with label "qjit" for both single and multi-qnode workflows. This helps eliminate any tech debt while still retaining user readability.

[mudit2812]

Looks good, just a few comments

[mehrdad2m]

Looks good on my side as well. Happy to approve once the open conversations are resolved. The only thing that scares me a little is the that the implementation is heavily dependent on the current IR structure which sometimes make it look hacky. That would be fine for MVP but we should also be careful not add too much tech debt.

[andrijapau]

The only thing that scares me a little is the that the implementation is heavily dependent on the current IR structure which sometimes make it look hacky.

Any section in particular? I would assume this one:

# If this is the jit_* FuncOp, only draw if there's more than one qnode (launch kernel)
        # This avoids redundant nested clusters: jit_my_circuit -> my_circuit -> ...
        visualize = True
        label = operation.sym_name.data
        if "jit_" in operation.sym_name.data:
            num_qnodes = 0
            for op in operation.body.ops:
                if isinstance(op, catalyst.LaunchKernelOp):
                    num_qnodes += 1
            # Get everything after the jit_* prefix
            label = str(label).split("_", maxsplit=1)[-1]
            if num_qnodes == 1:
                visualize = False

[andrijapau]

The only thing that scares me a little is the that the implementation is heavily dependent on the current IR structure which sometimes make it look hacky.

Any section in particular? I would assume this one:

# If this is the jit_* FuncOp, only draw if there's more than one qnode (launch kernel)
        # This avoids redundant nested clusters: jit_my_circuit -> my_circuit -> ...
        visualize = True
        label = operation.sym_name.data
        if "jit_" in operation.sym_name.data:
            num_qnodes = 0
            for op in operation.body.ops:
                if isinstance(op, catalyst.LaunchKernelOp):
                    num_qnodes += 1
            # Get everything after the jit_* prefix
            label = str(label).split("_", maxsplit=1)[-1]
            if num_qnodes == 1:
                visualize = False

I removed the hacky work around. See context here.

[andrijapau]

Suggested change

	# If it is a multi-qnode workflow, it will represent the "workflow" function
	# If it is a single qnode, it will represent the quantum function.