Merge branch 'main' into david/185-pyqrack-squin

Author: weinbe58

src/bloqade/pyqrack/base.py

@@ -26,13 +26,28 @@ class PyQrackOptions(typing.TypedDict):
     isOpenCL: bool
 
 
+def _validate_pyqrack_options(options: PyQrackOptions) -> None:
+    if options["isBinaryDecisionTree"] and options["isStabilizerHybrid"]:
+        raise ValueError(
+            "Cannot use both isBinaryDecisionTree and isStabilizerHybrid at the same time."
+        )
+    elif options["isTensorNetwork"] and options["isBinaryDecisionTree"]:
+        raise ValueError(
+            "Cannot use both isTensorNetwork and isBinaryDecisionTree at the same time."
+        )
+    elif options["isTensorNetwork"] and options["isStabilizerHybrid"]:
+        raise ValueError(
+            "Cannot use both isTensorNetwork and isStabilizerHybrid at the same time."
+        )
+
+
 def _default_pyqrack_args() -> PyQrackOptions:
     return PyQrackOptions(
         qubitCount=-1,
         isTensorNetwork=False,
         isSchmidtDecomposeMulti=True,
         isSchmidtDecompose=True,
-        isStabilizerHybrid=True,
+        isStabilizerHybrid=False,
         isBinaryDecisionTree=True,
         isPaged=True,
         isCpuGpuHybrid=True,
@@ -45,6 +60,9 @@ class MemoryABC(abc.ABC):
     pyqrack_options: PyQrackOptions = field(default_factory=_default_pyqrack_args)
     sim_reg: "QrackSimulator" = field(init=False)
 
+    def __post_init__(self):
+        _validate_pyqrack_options(self.pyqrack_options)
+
     @abc.abstractmethod
     def allocate(self, n_qubits: int) -> tuple[int, ...]:
         """Allocate `n_qubits` qubits and return their ids."""

src/bloqade/squin/op/stmts.py

@@ -148,14 +148,14 @@ class PauliOp(ConstantUnitary):
 
 
 @statement(dialect=dialect)
-class CliffordString(ConstantUnitary):
+class PauliString(ConstantUnitary):
     traits = frozenset({ir.Pure(), lowering.FromPythonCall(), Unitary(), HasSites()})
     string: str = info.attribute()
 
     def verify(self) -> None:
-        if not set("XYZHS").issuperset(self.string):
+        if not set("XYZ").issuperset(self.string):
             raise ValueError(
-                f"Invalid Clifford string: {self.string}. Must be a combination of 'X', 'Y', 'Z', 'H', and 'S'."
+                f"Invalid Pauli string: {self.string}. Must be a combination of 'X', 'Y', and 'Z'."
             )
 
 

src/bloqade/squin/wire.py

@@ -69,6 +69,25 @@ def __init__(self, operator: ir.SSAValue, *args: ir.SSAValue):
         )  # custom lowering required for wrapper to work here
 
 
+# Carry over from Qubit dialect
+@statement(dialect=dialect)
+class Broadcast(ir.Statement):
+    traits = frozenset({lowering.FromPythonCall(), ir.Pure()})
+    operator: ir.SSAValue = info.argument(OpType)
+    inputs: tuple[ir.SSAValue, ...] = info.argument(WireType)
+
+    def __init__(self, operator: ir.SSAValue, *args: ir.SSAValue):
+        result_types = tuple(WireType for _ in args)
+        super().__init__(
+            args=(operator,) + args,
+            result_types=result_types,
+            args_slice={
+                "operator": 0,
+                "inputs": slice(1, None),
+            },  # pretty printing + syntax sugar
+        )  # custom lowering required for wrapper to work here
+
+
 # NOTE: measurement cannot be pure because they will collapse the state
 #       of the qubit. The state is a hidden state that is not visible to
 #      the user in the wire dialect.
@@ -99,6 +118,7 @@ class Reset(ir.Statement):
 class ConstPropWire(interp.MethodTable):
 
     @interp.impl(Apply)
+    @interp.impl(Broadcast)
     def apply(self, interp, frame, stmt: Apply):
 
         return frame.get_values(stmt.inputs)

test/pyqrack/test_target.py

@@ -103,7 +103,9 @@ def multiple_registers():
 
         return q1
 
-    target = PyQrack(6)
+    target = PyQrack(
+        6, pyqrack_options={"isBinaryDecisionTree": False, "isStabilizerHybrid": True}
+    )
     q1 = target.run(multiple_registers)
 
     assert isinstance(q1, ilist.IList)

test/qasm2/passes/test_heuristic_noise.py

@@ -287,11 +287,11 @@ def test_method():
     expected_block = ir.Block(
         [
             n_qubits := constant.Constant(1),
-            reg0 := core.QRegNew(n_qubits.result),
-            zero := constant.Constant(0),
-            q0 := core.QRegGet(reg0.result, zero.result),
             reg1 := core.QRegNew(n_qubits.result),
+            zero := constant.Constant(0),
             q1 := core.QRegGet(reg1.result, zero.result),
+            reg0 := core.QRegNew(n_qubits.result),
+            q0 := core.QRegGet(reg0.result, zero.result),
             reg_list := ilist.New(
                 values=[reg0.result, reg1.result], elem_type=reg0.result.type
             ),

test/squin/test_constprop.py

@@ -0,0 +1,110 @@
+# There's a method table in the wire dialect statements
+# that handles the multiple return values from Apply and Broadcast
+# that can cause problems for constant propoagation's default implementation.
+# These tests just make sure there are corresponding lattice types per each
+# SSA value (as opposed to a bunch of "missing" entries despite multiple
+# return values from Broadcast and Apply)
+
+from kirin import ir, types
+from kirin.passes import Fold
+from kirin.dialects import py, func
+from kirin.analysis.const import Propagate
+
+from bloqade import squin
+
+
+def as_int(value: int):
+    return py.constant.Constant(value=value)
+
+
+def as_float(value: float):
+    return py.constant.Constant(value=value)
+
+
+def gen_func_from_stmts(stmts):
+
+    squin_with_py = squin.groups.wired.add(py)
+
+    block = ir.Block(stmts)
+    block.args.append_from(types.MethodType[[], types.NoneType], "main_self")
+    func_wrapper = func.Function(
+        sym_name="main",
+        signature=func.Signature(inputs=(), output=squin.op.types.OpType),
+        body=ir.Region(blocks=block),
+    )
+
+    constructed_method = ir.Method(
+        mod=None,
+        py_func=None,
+        sym_name="main",
+        dialects=squin_with_py,
+        code=func_wrapper,
+        arg_names=[],
+    )
+
+    fold_pass = Fold(squin_with_py)
+    fold_pass(constructed_method)
+
+    return constructed_method
+
+
+def test_wire_apply_constprop():
+
+    stmts = [
+        (n_qubits := as_int(2)),
+        (qreg := squin.qubit.New(n_qubits=n_qubits.result)),
+        (idx0 := as_int(0)),
+        (q0 := py.GetItem(qreg.result, idx0.result)),
+        (idx1 := as_int(1)),
+        (q1 := py.GetItem(qreg.result, idx1.result)),
+        # get wires
+        (w0 := squin.wire.Unwrap(q0.result)),
+        (w1 := squin.wire.Unwrap(q1.result)),
+        # put wire through gates
+        (x := squin.op.stmts.X()),
+        (cx := squin.op.stmts.Control(op=x.result, n_controls=1)),
+        (a := squin.wire.Apply(cx.result, w0.result, w1.result)),
+        (func.Return(a.results[0])),
+    ]
+    constructed_method = gen_func_from_stmts(stmts)
+
+    prop_analysis = Propagate(constructed_method.dialects)
+    frame, _ = prop_analysis.run_analysis(constructed_method)
+
+    assert len(frame.entries.values()) == 13
+
+
+def test_wire_broadcast_constprop():
+
+    stmts = [
+        (n_qubits := as_int(4)),
+        (qreg := squin.qubit.New(n_qubits=n_qubits.result)),
+        (idx0 := as_int(0)),
+        (q0 := py.GetItem(qreg.result, idx0.result)),
+        (idx1 := as_int(1)),
+        (q1 := py.GetItem(qreg.result, idx1.result)),
+        (idx2 := as_int(2)),
+        (q2 := py.GetItem(qreg.result, idx2.result)),
+        (idx3 := as_int(3)),
+        (q3 := py.GetItem(qreg.result, idx3.result)),
+        # get wires
+        (w0 := squin.wire.Unwrap(q0.result)),
+        (w1 := squin.wire.Unwrap(q1.result)),
+        (w2 := squin.wire.Unwrap(q2.result)),
+        (w3 := squin.wire.Unwrap(q3.result)),
+        # put wire through gates
+        (x := squin.op.stmts.X()),
+        (cx := squin.op.stmts.Control(op=x.result, n_controls=1)),
+        (
+            a := squin.wire.Broadcast(
+                cx.result, w0.result, w1.result, w2.result, w3.result
+            )
+        ),
+        (func.Return(a.results[0])),
+    ]
+    constructed_method = gen_func_from_stmts(stmts)
+
+    prop_analysis = Propagate(constructed_method.dialects)
+    frame, _ = prop_analysis.run_analysis(constructed_method)
+
+    assert len(frame.entries.values()) == 21