feat: add cluster visualization for FuncOp and DeviceInitOp operations

doc/releases/changelog-dev.md

@@ -7,6 +7,7 @@
   [(#2229)](https://github.com/PennyLaneAI/catalyst/pull/2229)
   [(#2214)](https://github.com/PennyLaneAI/catalyst/pull/2214)
   [(#2246)](https://github.com/PennyLaneAI/catalyst/pull/2246)
+  [(#2231)](https://github.com/PennyLaneAI/catalyst/pull/2231)
 
 * Added ``catalyst.switch``, a qjit compatible, index-switch style control flow decorator.
   [(#2171)](https://github.com/PennyLaneAI/catalyst/pull/2171)

frontend/catalyst/python_interface/visualization/construct_circuit_dag.py

@@ -16,9 +16,10 @@
 
 from functools import singledispatchmethod
 
-from xdsl.dialects import builtin
+from xdsl.dialects import builtin, func
 from xdsl.ir import Block, Operation, Region
 
+from catalyst.python_interface.dialects import catalyst, quantum
 from catalyst.python_interface.visualization.dag_builder import DAGBuilder
 
 
@@ -40,13 +41,27 @@ class ConstructCircuitDAG:
     def __init__(self, dag_builder: DAGBuilder) -> None:
         self.dag_builder: DAGBuilder = dag_builder
 
+        # Keep track of nesting clusters using a stack
+        self._cluster_uid_stack: list[str] = []
+
+        # Use counter internally for UID
+        self._node_uid_counter: int = 0
+        self._cluster_uid_counter: int = 0
+
+    def _reset(self) -> None:
+        """Resets the instance."""
+        self._cluster_uid_stack: list[str] = []
+        self._node_uid_counter: int = 0
+        self._cluster_uid_counter: int = 0
+
     def construct(self, module: builtin.ModuleOp) -> None:
         """Constructs the DAG from the module.
 
         Args:
             module (xdsl.builtin.ModuleOp): The module containing the quantum program to visualize.
 
         """
+        self._reset()
         for op in module.ops:
             self._visit_operation(op)
 
@@ -69,3 +84,57 @@ def _visit_block(self, block: Block) -> None:
         """Visit an xDSL Block operation, dispatching handling for each contained Operation."""
         for op in block.ops:
             self._visit_operation(op)
+
+    # ============
+    # DEVICE NODE
+    # ============
+
+    @_visit_operation.register
+    def _device_init(self, operation: quantum.DeviceInitOp) -> None:
+        """Handles the initialization of a quantum device."""
+        node_id = f"node{self._node_uid_counter}"
+        self.dag_builder.add_node(
+            node_id,
+            label=operation.device_name.data,
+            cluster_uid=self._cluster_uid_stack[-1],
+            fillcolor="grey",
+            color="black",
+            penwidth=2,
+            shape="rectangle",
+        )
+        self._node_uid_counter += 1
+
+    # =======================
+    # FuncOp NESTING UTILITY
+    # =======================
+
+    @_visit_operation.register
+    def _func_op(self, operation: func.FuncOp) -> None:
+        """Visit a FuncOp Operation."""
+
+        label = operation.sym_name.data
+        if "jit_" in operation.sym_name.data:
+            label = "qjit"
+
+        uid = f"cluster{self._cluster_uid_counter}"
+        parent_cluster_uid = None if self._cluster_uid_stack == [] else self._cluster_uid_stack[-1]
+        self.dag_builder.add_cluster(
+            uid,
+            label=label,
+            cluster_uid=parent_cluster_uid,
+        )
+        self._cluster_uid_counter += 1
+        self._cluster_uid_stack.append(uid)
+
+        self._visit_block(operation.regions[0].blocks[0])
+
+    @_visit_operation.register
+    def _func_return(self, operation: func.ReturnOp) -> None:
+        """Handle func.return to exit FuncOp's cluster scope."""
+
+        # NOTE: Skip first cluster as it is the "base" of the graph diagram.
+        # In our case, it is the `qjit` bounding box.
+        if len(self._cluster_uid_stack) > 1:
+            # If we hit a func.return operation we know we are leaving
+            # the FuncOp's scope and so we can pop the ID off the stack.
+            self._cluster_uid_stack.pop()

frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py

@@ -19,12 +19,14 @@
 
 pytestmark = pytest.mark.usefixtures("requires_xdsl")
 
+# pylint: disable=wrong-import-position
+# This import needs to be after pytest in order to prevent ImportErrors
+import pennylane as qml
 from xdsl.dialects import test
 from xdsl.dialects.builtin import ModuleOp
 from xdsl.ir.core import Block, Region
 
-# pylint: disable=wrong-import-position
-# This import needs to be after pytest in order to prevent ImportErrors
+from catalyst.python_interface.conversion import xdsl_from_qjit
 from catalyst.python_interface.visualization.construct_circuit_dag import (
     ConstructCircuitDAG,
 )
@@ -124,3 +126,164 @@ def test_does_not_mutate_module():
 
     # Ensure not mutated
     assert str(module_op) == module_op_str_before
+
+
+@pytest.mark.unit
+class TestFuncOpVisualization:
+    """Tests the visualization of FuncOps with bounding boxes"""
+
+    def test_standard_qnode(self):
+        """Tests that a standard QJIT'd QNode is visualized correctly"""
+        dev = qml.device("null.qubit", wires=1)
+
+        @xdsl_from_qjit
+        @qml.qjit(autograph=True, target="mlir")
+        @qml.qnode(dev)
+        def my_workflow():
+            qml.H(0)
+
+        module = my_workflow()
+
+        utility = ConstructCircuitDAG(FakeDAGBuilder())
+        utility.construct(module)
+
+        graph_clusters = utility.dag_builder.clusters
+
+        # Check nesting is correct
+        # graph
+        # └── qjit
+        #     └── my_workflow
+
+        # Check qjit is nested under graph
+        assert graph_clusters["cluster0"]["cluster_label"] == "qjit"
+        assert graph_clusters["cluster0"]["parent_cluster_uid"] is None
+
+        # Check that my_workflow is under qjit
+        assert graph_clusters["cluster1"]["cluster_label"] == "my_workflow"
+        assert graph_clusters["cluster1"]["parent_cluster_uid"] == "cluster0"
+
+    def test_nested_qnodes(self):
+        """Tests that nested QJIT'd QNodes are visualized correctly"""
+
+        dev = qml.device("null.qubit", wires=1)
+
+        @qml.qnode(dev)
+        def my_qnode2():
+            qml.X(0)
+
+        @qml.qnode(dev)
+        def my_qnode1():
+            qml.H(0)
+
+        @xdsl_from_qjit
+        @qml.qjit(autograph=True, target="mlir")
+        def my_workflow():
+            my_qnode1()
+            my_qnode2()
+
+        module = my_workflow()
+
+        utility = ConstructCircuitDAG(FakeDAGBuilder())
+        utility.construct(module)
+
+        graph_clusters = utility.dag_builder.clusters
+
+        # Check nesting is correct
+        # graph
+        # └── qjit
+        #     ├── my_qnode1
+        #     └── my_qnode2
+
+        # Check qjit is under graph
+        assert graph_clusters["cluster0"]["cluster_label"] == "qjit"
+        assert graph_clusters["cluster0"]["parent_cluster_uid"] is None
+
+        # Check both qnodes are under my_workflow
+        assert graph_clusters["cluster1"]["cluster_label"] == "my_qnode1"
+        assert graph_clusters["cluster1"]["parent_cluster_uid"] == "cluster0"
+
+        assert graph_clusters["cluster2"]["cluster_label"] == "my_qnode2"
+        assert graph_clusters["cluster2"]["parent_cluster_uid"] == "cluster0"
+
+
+class TestDeviceNode:
+    """Tests that the device node is correctly visualized."""
+
+    def test_standard_qnode(self):
+        """Tests that a standard setup works."""
+
+        dev = qml.device("null.qubit", wires=1)
+
+        @xdsl_from_qjit
+        @qml.qjit(autograph=True, target="mlir")
+        @qml.qnode(dev)
+        def my_workflow():
+            qml.H(0)
+
+        module = my_workflow()
+
+        utility = ConstructCircuitDAG(FakeDAGBuilder())
+        utility.construct(module)
+
+        graph_nodes = utility.dag_builder.nodes
+        graph_clusters = utility.dag_builder.clusters
+
+        # Check nesting is correct
+        # graph
+        # └── qjit
+        #     └── my_workflow: NullQubit
+
+        # Assert device node is inside my_workflow cluster
+        assert graph_clusters["cluster1"]["cluster_label"] == "my_workflow"
+        assert graph_nodes["node0"]["parent_cluster_uid"] == "cluster1"
+
+        # Assert label is as expected
+        assert graph_nodes["node0"]["label"] == "NullQubit"
+
+    def test_nested_qnodes(self):
+        """Tests that nested QJIT'd QNodes are visualized correctly"""
+
+        dev1 = qml.device("null.qubit", wires=1)
+        dev2 = qml.device("lightning.qubit", wires=1)
+
+        @qml.qnode(dev2)
+        def my_qnode2():
+            qml.X(0)
+
+        @qml.qnode(dev1)
+        def my_qnode1():
+            qml.H(0)
+
+        @xdsl_from_qjit
+        @qml.qjit(autograph=True, target="mlir")
+        def my_workflow():
+            my_qnode1()
+            my_qnode2()
+
+        module = my_workflow()
+
+        utility = ConstructCircuitDAG(FakeDAGBuilder())
+        utility.construct(module)
+
+        graph_nodes = utility.dag_builder.nodes
+        graph_clusters = utility.dag_builder.clusters
+
+        # Check nesting is correct
+        # graph
+        # └── qjit
+        #     ├── my_qnode1: NullQubit
+        #     └── my_qnode2: LightningSimulator
+
+        # Assert lightning.qubit device node is inside my_qnode1 cluster
+        assert graph_clusters["cluster1"]["cluster_label"] == "my_qnode1"
+        assert graph_nodes["node0"]["parent_cluster_uid"] == "cluster1"
+
+        # Assert label is as expected
+        assert graph_nodes["node0"]["label"] == "NullQubit"
+
+        # Assert null qubit device node is inside my_qnode2 cluster
+        assert graph_clusters["cluster2"]["cluster_label"] == "my_qnode2"
+        assert graph_nodes["node1"]["parent_cluster_uid"] == "cluster2"
+
+        # Assert label is as expected
+        assert graph_nodes["node1"]["label"] == "LightningSimulator"