Add a graph_state_prep operation to the xDSL MBQC dialect (#8059)

**Context:** We recently added the `mbqc.graph_state_prep` operation to
Catalyst in PennyLaneAI/catalyst#1965. We now
need the equivalent operation defined in the MBQC dialect of the unified
Python compiler.

**Description of the Change:** Adds the `mbqc.graph_state_prep`
operation to the xDSL MBQC dialect. It's implementation is equivalent to
the one already defined in Catalyst.

[sc-97308]

Author: joeycarter

doc/releases/changelog-dev.md

@@ -192,6 +192,7 @@
 * The `mbqc` xDSL dialect has been added to the Python compiler, which is used to represent
   measurement-based quantum-computing instructions in the xDSL framework.
   [(#7815)](https://github.com/PennyLaneAI/pennylane/pull/7815)
+  [(#8059)](https://github.com/PennyLaneAI/pennylane/pull/8059)
 
 * The `AllocQubitOp` and `DeallocQubitOp` operations have been added to the `Quantum` dialect in the
   Python compiler.

pennylane/compiler/python_compiler/dialects/mbqc.py

@@ -19,11 +19,24 @@
 
 It was initially generated by xDSL (using the ``xdsl-tblgen`` tool) starting from the
 catalyst/mlir/include/MBQC/IR/MBQCDialect.td file in the catalyst repository.
+
+For detailed documentation on the operations contained in this dialect, please refer to the MBQC
+dialect documentation in Catalyst.
 """
 
 from typing import TypeAlias
 
-from xdsl.dialects.builtin import I32, Float64Type, IntegerAttr, IntegerType
+from xdsl.dialects.builtin import (
+    I32,
+    AnyAttr,
+    Float64Type,
+    IntegerAttr,
+    IntegerType,
+    MemRefType,
+    StringAttr,
+    TensorType,
+    i1,
+)
 from xdsl.ir import Dialect, EnumAttribute, Operation, SpacedOpaqueSyntaxAttribute, SSAValue
 from xdsl.irdl import (
     IRDLOperation,
@@ -37,7 +50,8 @@
 from xdsl.utils.exceptions import VerifyException
 from xdsl.utils.str_enum import StrEnum  # StrEnum is standard in Python>=3.11
 
-from .quantum import QubitType
+from ..xdsl_extras import MemRefRankConstraint, TensorRankConstraint
+from .quantum import QubitType, QuregType
 
 QubitSSAValue: TypeAlias = SSAValue[QubitType]
 
@@ -69,6 +83,10 @@ class MeasureInBasisOp(IRDLOperation):
 
     name = "mbqc.measure_in_basis"
 
+    assembly_format = """
+            `[` $plane `,` $angle `]` $in_qubit (`postselect` $postselect^)? attr-dict `:` type(results)
+        """
+
     in_qubit = operand_def(QubitType)
 
     plane = prop_def(MeasurementPlaneAttr)
@@ -81,10 +99,6 @@ class MeasureInBasisOp(IRDLOperation):
 
     out_qubit = result_def(QubitType)
 
-    assembly_format = """
-            `[` $plane `,` $angle `]` $in_qubit (`postselect` $postselect^)? attr-dict `:` type(results)
-        """
-
     def __init__(
         self,
         in_qubit: QubitSSAValue | Operation,
@@ -115,10 +129,54 @@ def verify_(self):
             raise VerifyException("'postselect' must be 0 or 1.")
 
 
+@irdl_op_definition
+class GraphStatePrepOp(IRDLOperation):
+    """Allocate resources for a new graph state."""
+
+    # pylint: disable=too-few-public-methods
+
+    name = "mbqc.graph_state_prep"
+
+    assembly_format = """
+            `(` $adj_matrix `:` type($adj_matrix) `)` `[` `init` $init_op `,` `entangle` $entangle_op `]` attr-dict `:` type(results)
+        """
+
+    adj_matrix = operand_def(
+        (TensorType.constr(i1) & TensorRankConstraint(1))
+        | (MemRefType.constr(i1) & MemRefRankConstraint(1))
+    )
+
+    init_op = prop_def(StringAttr)
+
+    entangle_op = prop_def(StringAttr)
+
+    qreg = result_def(QuregType)
+
+    def __init__(
+        self, adj_matrix: AnyAttr, init_op: str | StringAttr, entangle_op: str | StringAttr
+    ):
+        if isinstance(init_op, str):
+            init_op = StringAttr(data=init_op)
+
+        if isinstance(entangle_op, str):
+            entangle_op = StringAttr(data=entangle_op)
+
+        properties = {"init_op": init_op, "entangle_op": entangle_op}
+
+        qreg = QuregType()
+
+        super().__init__(
+            operands=(adj_matrix,),
+            result_types=(qreg,),
+            properties=properties,
+        )
+
+
 MBQC = Dialect(
     "mbqc",
     [
         MeasureInBasisOp,
+        GraphStatePrepOp,
     ],
     [
         MeasurementPlaneAttr,

tests/python_compiler/dialects/test_mbqc_dialect.py

@@ -23,7 +23,7 @@
 
 pytestmark = pytest.mark.external
 
-from xdsl.dialects import builtin, test
+from xdsl.dialects import arith, builtin, test
 from xdsl.utils.exceptions import VerifyException
 
 from pennylane.compiler.python_compiler.dialects import Quantum, mbqc
@@ -33,6 +33,7 @@
 
 expected_ops_names = {
     "MeasureInBasisOp": "mbqc.measure_in_basis",
+    "GraphStatePrepOp": "mbqc.graph_state_prep",
 }
 
 expected_attrs_names = {
@@ -90,6 +91,11 @@ def test_assembly_format(run_filecheck):
     // COM: Check generic format
     // CHECK: {{%.+}}, {{%.+}} = mbqc.measure_in_basis[XY, [[angle]]] [[qubit]] postselect 0 : i1, !quantum.bit
     %res:2 = "mbqc.measure_in_basis"(%qubit, %angle) <{plane = #mbqc<measurement_plane XY>, postselect = 0 : i32}> : (!quantum.bit, f64) -> (i1, !quantum.bit)
+
+    // CHECK: [[adj_matrix:%.+]] = arith.constant {{.*}} : tensor<6xi1>
+    // CHECK: [[graph_reg:%.+]] = mbqc.graph_state_prep{{\s*}}([[adj_matrix]] : tensor<6xi1>) [init "Hadamard", entangle "CZ"] : !quantum.reg
+    %adj_matrix = arith.constant dense<[1, 0, 1, 0, 0, 1]> : tensor<6xi1>
+    %graph_reg = mbqc.graph_state_prep (%adj_matrix : tensor<6xi1>) [init "Hadamard", entangle "CZ"] : !quantum.reg
     """
 
     run_filecheck(program)
@@ -154,3 +160,19 @@ def test_invalid_postselect_raises_on_verify(self, postselect):
 
         with pytest.raises(VerifyException, match="'postselect' must be 0 or 1"):
             measure_in_basis_op.verify_()
+
+    @pytest.mark.parametrize("init_op", ["Hadamard", builtin.StringAttr(data="Hadamard")])
+    @pytest.mark.parametrize("entangle_op", ["CZ", builtin.StringAttr(data="CZ")])
+    def test_graph_state_prep_instantiation(self, init_op, entangle_op):
+        """Test the instantiation of a mbqc.graph_state_prep op."""
+        adj_matrix = [1, 0, 1, 0, 0, 1]
+        adj_matrix_op = arith.ConstantOp(
+            builtin.DenseIntOrFPElementsAttr.from_list(
+                type=builtin.TensorType(builtin.IntegerType(1), shape=(6,)), data=adj_matrix
+            )
+        )
+        graph_state_prep_op = mbqc.GraphStatePrepOp(adj_matrix_op.result, init_op, entangle_op)
+
+        assert graph_state_prep_op.adj_matrix == adj_matrix_op.result
+        assert graph_state_prep_op.init_op == builtin.StringAttr(data="Hadamard")
+        assert graph_state_prep_op.entangle_op == builtin.StringAttr(data="CZ")