Fancy constraints for qssa

inconspiquous/alloc.py

@@ -1,6 +1,8 @@
 from abc import ABC
-from xdsl.ir import ParametrizedAttribute
+from typing import Sequence
+from xdsl.ir import Attribute, ParametrizedAttribute
+from xdsl.irdl import WithRangeType
 
 
-class AllocAttr(ParametrizedAttribute, ABC):
-    pass
+class AllocAttr(ParametrizedAttribute, WithRangeType, ABC):
+    def get_types(self) -> Sequence[Attribute]: ...

inconspiquous/dialects/qssa.py

@@ -1,6 +1,6 @@
 from typing import ClassVar
 from xdsl.dialects.builtin import i1
-from xdsl.ir import Dialect, Operation, SSAValue, VerifyException
+from xdsl.ir import Dialect, Operation, SSAValue
 from xdsl.irdl import (
     EqAttrConstraint,
     IRDLOperation,
@@ -15,23 +15,23 @@
     var_result_def,
 )
 
-from inconspiquous.gates import GateAttr
+from inconspiquous.gates import GateAttr, GateConstraint
 from inconspiquous.dialects.qubit import BitType
 
 
 @irdl_op_definition
 class GateOp(IRDLOperation):
     name = "qssa.gate"
 
-    gate = prop_def(GateAttr)
-
-    _T: ClassVar[RangeConstraint] = RangeVarConstraint(
-        "T", RangeOf(EqAttrConstraint(BitType()))
+    _Q: ClassVar[RangeConstraint] = RangeVarConstraint(
+        "Q", RangeOf(EqAttrConstraint(BitType()))
     )
 
-    ins = var_operand_def(_T)
+    gate = prop_def(GateConstraint(_Q))
+
+    ins = var_operand_def(_Q)
 
-    outs = var_result_def(_T)
+    outs = var_result_def(_Q)
 
     assembly_format = "`<` $gate `>` $ins attr-dict `:` type($ins)"
 
@@ -44,12 +44,6 @@ def __init__(self, gate: GateAttr, *ins: SSAValue | Operation):
             result_types=tuple(BitType() for _ in ins),
         )
 
-    def verify_(self) -> None:
-        if len(self.ins) != self.gate.num_qubits:
-            raise VerifyException(
-                f"Gate {self.gate} expected {self.gate.num_qubits} input qubits but got {len(self.ins)}."
-            )
-
 
 @irdl_op_definition
 class MeasureOp(IRDLOperation):

inconspiquous/dialects/qubit.py

@@ -1,10 +1,17 @@
-from xdsl.ir import Dialect, ParametrizedAttribute, TypeAttribute
+from typing import ClassVar, Sequence
+from xdsl.ir import Attribute, Dialect, ParametrizedAttribute, TypeAttribute
 from xdsl.irdl import (
+    AnyAttr,
+    BaseAttr,
     IRDLOperation,
+    RangeConstraint,
+    RangeOf,
+    RangeVarConstraint,
+    WithRangeTypeConstraint,
     irdl_attr_definition,
     irdl_op_definition,
     prop_def,
-    result_def,
+    var_result_def,
 )
 
 from inconspiquous.alloc import AllocAttr
@@ -27,14 +34,21 @@ class AllocZeroAttr(AllocAttr):
 
     name = "qubit.zero"
 
+    def get_types(self) -> Sequence[Attribute]:
+        return (BitType(),)
+
 
 @irdl_op_definition
 class AllocOp(IRDLOperation):
     name = "qubit.alloc"
 
-    alloc = prop_def(AllocAttr, default_value=AllocZeroAttr())
+    _T: ClassVar[RangeConstraint] = RangeVarConstraint("T", RangeOf(AnyAttr()))
+
+    alloc = prop_def(
+        WithRangeTypeConstraint(BaseAttr(AllocAttr), _T), default_value=AllocZeroAttr()
+    )
 
-    outs = result_def(BitType())
+    outs = var_result_def(_T)
 
     assembly_format = "(`` `<` $alloc^ `>`)? attr-dict"
 
@@ -43,7 +57,7 @@ def __init__(self, alloc: AllocAttr = AllocZeroAttr()):
             properties={
                 "alloc": alloc,
             },
-            result_types=[BitType()],
+            result_types=[alloc.get_types()],
         )
 
 

inconspiquous/gates.py

@@ -1,7 +1,15 @@
 from abc import ABC
 
 from xdsl.ir import (
+    Attribute,
     ParametrizedAttribute,
+    VerifyException,
+)
+from xdsl.irdl import (
+    BaseAttr,
+    ConstraintContext,
+    GenericAttrConstraint,
+    RangeVarConstraint,
 )
 
 
@@ -33,3 +41,32 @@ class TwoQubitGate(GateAttr):
     @property
     def num_qubits(self) -> int:
         return 2
+
+
+class GateConstraint(GenericAttrConstraint[GateAttr]):
+    """
+    Constrains a given range variable to have the correct size for the gate.
+    """
+
+    range_var: str
+    gate_constraint: GenericAttrConstraint[GateAttr]
+
+    def __init__(
+        self,
+        range_constraint: str | RangeVarConstraint[Attribute],
+        gate_constraint: GenericAttrConstraint[GateAttr] = BaseAttr[GateAttr](GateAttr),
+    ):
+        if isinstance(range_constraint, str):
+            self.range_var = range_constraint
+        else:
+            self.range_var = range_constraint.name
+        self.gate_constraint = gate_constraint
+
+    def verify(self, attr: Attribute, constraint_context: ConstraintContext) -> None:
+        if self.range_var in constraint_context.range_variables:
+            attrs = constraint_context.get_range_variable(self.range_var)
+            assert isinstance(attr, GateAttr)
+            if attr.num_qubits != len(attrs):
+                raise VerifyException(
+                    f"Gate {attr} expected {attr.num_qubits} qubits but got {len(attrs)}"
+                )

pyproject.toml

@@ -20,7 +20,7 @@ dev-dependencies = [
 ]
 
 [tool.uv.sources]
-xdsl = { git = "https://github.com/xdslproject/xdsl", branch = "alexarice/range-variable-sequence" }
+xdsl = { git = "https://github.com/xdslproject/xdsl", branch = "alexarice/constraint-branch" }
 
 [project.scripts]
 quopt = "inconspiquous.tools.quopt_main:main"

tests/filecheck/dialects/qssa/gate_counts.mlir

@@ -2,20 +2,20 @@
 
 %q0 = qubit.alloc
 
-// CHECK: attributes ('!qubit.bit',) expected from range variable 'T', but got ('!qubit.bit', '!qubit.bit')
+// CHECK: attributes ('!qubit.bit',) expected from range variable 'Q', but got ('!qubit.bit', '!qubit.bit')
 %q1, %q2 = "qssa.gate"(%q0) <{"gate" = #gate.cnot}> : (!qubit.bit) -> (!qubit.bit, !qubit.bit)
 
 // -----
 
 %q0 = qubit.alloc
 %q1 = qubit.alloc
 
-// CHECK: attributes ('!qubit.bit', '!qubit.bit') expected from range variable 'T', but got ('!qubit.bit',)
+// CHECK: attributes ('!qubit.bit', '!qubit.bit') expected from range variable 'Q', but got ('!qubit.bit',)
 %q2 = "qssa.gate"(%q0, %q1) <{"gate" = #gate.cnot}> : (!qubit.bit, !qubit.bit) -> !qubit.bit
 
 // -----
 
 %q0 = qubit.alloc
 
-// CHECK: Gate #gate.cnot expected 2 input qubits but got 1
+// CHECK: Gate #gate.cnot expected 2 qubits but got 1
 %q1 = "qssa.gate"(%q0) <{"gate" = #gate.cnot}> : (!qubit.bit) -> !qubit.bit