feat: add cluster visualization for WhileOp, IfOp, ForOp

doc/releases/changelog-dev.md

@@ -7,6 +7,8 @@
   [(#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)
+  [(#2234)](https://github.com/PennyLaneAI/catalyst/pull/2234)
 
 * 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,11 @@
 
 from functools import singledispatchmethod
 
-from xdsl.dialects import builtin
-from xdsl.ir import Block, Operation, Region
+from xdsl.dialects import builtin, func, scf
+from xdsl.ir import Block, Operation, Region, SSAValue
 
+from catalyst.python_interface.dialects import catalyst, quantum
+from catalyst.python_interface.inspection.xdsl_conversion import resolve_constant_params
 from catalyst.python_interface.visualization.dag_builder import DAGBuilder
 
 
@@ -40,13 +42,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 +85,179 @@ 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)
+
+    # =============
+    # CONTROL FLOW
+    # =============
+
+    @_visit_operation.register
+    def _for_op(self, operation: scf.ForOp) -> None:
+        """Handle an xDSL ForOp operation."""
+
+        # TODO: Extract from IR in future PR
+        iter_var = "..."
+        start, stop, step = "...", "...", "..."
+        label = f"for {iter_var} in range({start}, {stop}, {step})"
+
+        uid = f"cluster{self._cluster_uid_counter}"
+        self.dag_builder.add_cluster(
+            uid,
+            node_label=label,
+            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 _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 ...",
+            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",
+            fontsize=10,
+            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):
+
+            def _get_conditional_branch_label(i):
+                if i == 0:
+                    return "if ..."
+                elif i == num_regions - 1:
+                    return "else"
+                else:
+                    return "elif ..."
+
+            uid = f"cluster{self._cluster_uid_counter}"
+            self.dag_builder.add_cluster(
+                uid,
+                node_label=_get_conditional_branch_label(i),
+                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
+    # ============
+
+    @_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()
+
+
+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