feat: visualize static operators and measurements as nodes

doc/releases/changelog-dev.md

@@ -8,6 +8,8 @@
   [(#2214)](https://github.com/PennyLaneAI/catalyst/pull/2214)
   [(#2246)](https://github.com/PennyLaneAI/catalyst/pull/2246)
   [(#2231)](https://github.com/PennyLaneAI/catalyst/pull/2231)
+  [(#2234)](https://github.com/PennyLaneAI/catalyst/pull/2234)
+  [(#2118)](https://github.com/PennyLaneAI/catalyst/pull/2218)
 
 * 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

@@ -14,12 +14,18 @@
 
 """Contains the ConstructCircuitDAG tool for constructing a DAG from an xDSL module."""
 
-from functools import singledispatchmethod
-
-from xdsl.dialects import builtin, func
-from xdsl.ir import Block, Operation, Region
-
-from catalyst.python_interface.dialects import catalyst, quantum
+from functools import singledispatch, singledispatchmethod
+
+from pennylane.measurements import MeasurementProcess
+from pennylane.operation import Operator
+from xdsl.dialects import builtin, func, scf
+from xdsl.ir import Block, Operation, Region, SSAValue
+
+from catalyst.python_interface.dialects import quantum
+from catalyst.python_interface.inspection.xdsl_conversion import (
+    xdsl_to_qml_measurement,
+    xdsl_to_qml_op,
+)
 from catalyst.python_interface.visualization.dag_builder import DAGBuilder
 
 
@@ -85,6 +91,206 @@ def _visit_block(self, block: Block) -> None:
         for op in block.ops:
             self._visit_operation(op)
 
+    # ===================
+    # QUANTUM OPERATIONS
+    # ===================
+
+    @_visit_operation.register
+    def _unitary(
+        self,
+        op: quantum.CustomOp | quantum.GlobalPhaseOp | quantum.QubitUnitaryOp | quantum.MultiRZOp,
+    ) -> None:
+        """Generic handler for unitary gates."""
+
+        # Create PennyLane instance
+        qml_op = xdsl_to_qml_op(op)
+
+        # Add node to current cluster
+        node_uid = f"node{self._node_uid_counter}"
+        self.dag_builder.add_node(
+            uid=node_uid,
+            label=get_label(qml_op),
+            cluster_uid=self._cluster_uid_stack[-1],
+            # NOTE: "record" allows us to use ports (https://graphviz.org/doc/info/shapes.html#record)
+            shape="record",
+        )
+        self._node_uid_counter += 1
+
+    @_visit_operation.register
+    def _projective_measure_op(self, op: quantum.MeasureOp) -> None:
+        """Handler for the single-qubit projective measurement operation."""
+
+        # Create PennyLane instance
+        meas = xdsl_to_qml_measurement(op)
+
+        # Add node to current cluster
+        node_uid = f"node{self._node_uid_counter}"
+        self.dag_builder.add_node(
+            uid=node_uid,
+            label=get_label(meas),
+            cluster_uid=self._cluster_uid_stack[-1],
+            # NOTE: "record" allows us to use ports (https://graphviz.org/doc/info/shapes.html#record)
+            shape="record",
+        )
+        self._node_uid_counter += 1
+
+    # =====================
+    # QUANTUM MEASUREMENTS
+    # =====================
+
+    @_visit_operation.register
+    def _state_op(self, op: quantum.StateOp) -> None:
+        """Handler for the terminal state measurement operation."""
+
+        # Create PennyLane instance
+        meas = xdsl_to_qml_measurement(op)
+
+        # Add node to current cluster
+        node_uid = f"node{self._node_uid_counter}"
+        self.dag_builder.add_node(
+            uid=node_uid,
+            label=get_label(meas),
+            cluster_uid=self._cluster_uid_stack[-1],
+            fillcolor="lightpink",
+            color="lightpink3",
+        )
+        self._node_uid_counter += 1
+
+    @_visit_operation.register
+    def _statistical_measurement_ops(
+        self,
+        op: quantum.ExpvalOp | quantum.VarianceOp,
+    ) -> None:
+        """Handler for statistical measurement operations."""
+
+        # Create PennyLane instance
+        obs_op = op.obs.owner
+        meas = xdsl_to_qml_measurement(op, xdsl_to_qml_measurement(obs_op))
+
+        # Add node to current cluster
+        node_uid = f"node{self._node_uid_counter}"
+        self.dag_builder.add_node(
+            uid=node_uid,
+            label=get_label(meas),
+            cluster_uid=self._cluster_uid_stack[-1],
+            fillcolor="lightpink",
+            color="lightpink3",
+        )
+        self._node_uid_counter += 1
+
+    @_visit_operation.register
+    def _visit_sample_and_probs_ops(
+        self,
+        op: quantum.SampleOp | quantum.ProbsOp,
+    ) -> None:
+        """Handler for sample operations."""
+
+        # Create PennyLane instance
+        obs_op = op.obs.owner
+
+        # TODO: This doesn't logically make sense, but quantum.compbasis
+        # is obs_op and function below just pulls out the static wires
+        wires = xdsl_to_qml_measurement(obs_op)
+        meas = xdsl_to_qml_measurement(op, wires=None if wires == [] else wires)
+
+        # Add node to current cluster
+        node_uid = f"node{self._node_uid_counter}"
+        self.dag_builder.add_node(
+            uid=node_uid,
+            label=get_label(meas),
+            cluster_uid=self._cluster_uid_stack[-1],
+            fillcolor="lightpink",
+            color="lightpink3",
+        )
+        self._node_uid_counter += 1
+
+    # =============
+    # CONTROL FLOW
+    # =============
+
+    @_visit_operation.register
+    def _for_op(self, operation: scf.ForOp) -> None:
+        """Handle an xDSL ForOp operation."""
+
+        uid = f"cluster{self._cluster_uid_counter}"
+        self.dag_builder.add_cluster(
+            uid,
+            node_label="for loop",
+            label="",
+            cluster_uid=self._cluster_uid_stack[-1],
+        )
+        self._cluster_uid_stack.append(uid)
+        self._cluster_uid_counter += 1
+
+        self._visit_region(operation.regions[0])
+
+        self._cluster_uid_stack.pop()
+
+    @_visit_operation.register
+    def _while_op(self, operation: scf.WhileOp) -> None:
+        """Handle an xDSL WhileOp operation."""
+        uid = f"cluster{self._cluster_uid_counter}"
+        self.dag_builder.add_cluster(
+            uid,
+            node_label="while loop",
+            label="",
+            cluster_uid=self._cluster_uid_stack[-1],
+        )
+        self._cluster_uid_stack.append(uid)
+        self._cluster_uid_counter += 1
+
+        for region in operation.regions:
+            self._visit_region(region)
+
+        self._cluster_uid_stack.pop()
+
+    @_visit_operation.register
+    def _if_op(self, operation: scf.IfOp):
+        """Handles the scf.IfOp operation."""
+        flattened_if_op: list[tuple[SSAValue | None, Region]] = _flatten_if_op(operation)
+
+        uid = f"cluster{self._cluster_uid_counter}"
+        self.dag_builder.add_cluster(
+            uid,
+            node_label="",
+            label="conditional",
+            labeljust="l",
+            cluster_uid=self._cluster_uid_stack[-1],
+        )
+        self._cluster_uid_stack.append(uid)
+        self._cluster_uid_counter += 1
+
+        # Loop through each branch and visualize as a cluster
+        num_regions = len(flattened_if_op)
+        for i, (condition_ssa, region) in enumerate(flattened_if_op):
+            node_label = "elif"
+            if i == 0:
+                node_label = "if"
+            elif i == num_regions - 1:
+                node_label = "else"
+
+            uid = f"cluster{self._cluster_uid_counter}"
+            self.dag_builder.add_cluster(
+                uid,
+                node_label=node_label,
+                label="",
+                style="dashed",
+                penwidth=1,
+                cluster_uid=self._cluster_uid_stack[-1],
+            )
+            self._cluster_uid_stack.append(uid)
+            self._cluster_uid_counter += 1
+
+            # Go recursively into the branch to process internals
+            self._visit_region(region)
+
+            # Pop branch cluster after processing to ensure
+            # logical branches are treated as 'parallel'
+            self._cluster_uid_stack.pop()
+
+        # Pop IfOp cluster before leaving this handler
+        self._cluster_uid_stack.pop()
+
     # ============
     # DEVICE NODE
     # ============
@@ -138,3 +344,51 @@ def _func_return(self, operation: func.ReturnOp) -> None:
             # 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()
+
+
+def _flatten_if_op(op: scf.IfOp) -> list[tuple[SSAValue | None, Region]]:
+    """Recursively flattens a nested IfOp (if/elif/else chains)."""
+
+    condition_ssa: SSAValue = op.operands[0]
+    then_region, else_region = op.regions
+
+    # Save condition SSA in case we want to visualize it eventually
+    flattened_op: list[tuple[SSAValue | None, Region]] = [(condition_ssa, then_region)]
+
+    # Peak into else region to see if there's another IfOp
+    else_block: Block = else_region.block
+    # Completely relies on the structure that the second last operation
+    # will be an IfOp (seems to hold true)
+    if isinstance(else_block.ops.last.prev_op, scf.IfOp):
+        # Recursively flatten any IfOps found in said block
+        nested_flattened_op = _flatten_if_op(else_block.ops.last.prev_op)
+        flattened_op.extend(nested_flattened_op)
+        return flattened_op
+
+    # No more nested IfOps, therefore append final region
+    # with no SSAValue
+    flattened_op.extend([(None, else_region)])
+    return flattened_op
+
+
+@singledispatch
+def get_label(op: Operator | MeasurementProcess) -> str:
+    """Gets the appropriate label for a PennyLane object."""
+    return str(op)
+
+
+@get_label.register
+def _operator(op: Operator) -> str:
+    """Returns the appropriate label for an xDSL operation."""
+    wires = list(op.wires.labels)
+    if wires == []:
+        wires_str = "all"
+    else:
+        wires_str = f"[{', '.join(map(str, wires))}]"
+    return f"<name> {op.name}|<wire> {wires_str}"
+
+
+@get_label.register
+def _mp(mp: MeasurementProcess) -> str:
+    """Returns the appropriate label for an xDSL operation."""
+    return str(mp)