Add Broadcast to wire dialect

src/bloqade/squin/qubit.py

@@ -20,13 +20,6 @@
 dialect = ir.Dialect("squin.qubit")
 
 
-@statement(dialect=dialect)
-class New(ir.Statement):
-    traits = frozenset({lowering.FromPythonCall()})
-    n_qubits: ir.SSAValue = info.argument(types.Int)
-    result: ir.ResultValue = info.result(ilist.IListType[QubitType, types.Any])
-
-
 @statement(dialect=dialect)
 class Apply(ir.Statement):
     traits = frozenset({lowering.FromPythonCall()})
@@ -42,6 +35,12 @@ class Broadcast(ir.Statement):
 
 
 @statement(dialect=dialect)
+class New(ir.Statement):
+    traits = frozenset({lowering.FromPythonCall()})
+    n_qubits: ir.SSAValue = info.argument(types.Int)
+    result: ir.ResultValue = info.result(ilist.IListType[QubitType, types.Any])
+
+
 class MeasureAny(ir.Statement):
     name = "measure"
 
@@ -62,7 +61,6 @@ class MeasureQubit(ir.Statement):
 @statement(dialect=dialect)
 class MeasureQubitList(ir.Statement):
     name = "measure.qubit.list"
-
     traits = frozenset({lowering.FromPythonCall()})
     qubits: ir.SSAValue = info.argument(ilist.IListType[QubitType])
     result: ir.ResultValue = info.result(ilist.IListType[types.Bool])

src/bloqade/squin/wire.py

@@ -51,7 +51,25 @@ class Unwrap(ir.Statement):
 # In Quake, you put a wire in and get a wire out when you "apply" an operator
 # In this case though we just need to indicate that an operator is applied to list[wires]
 @statement(dialect=dialect)
-class Apply(ir.Statement):  # apply(op, w1, w2, ...)
+class Apply(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,  # result types of the Apply statement, should all be WireTypes
+            args_slice={
+                "operator": 0,
+                "inputs": slice(1, None),
+            },  # pretty printing + syntax sugar
+        )  # custom lowering required for wrapper to work here
+
+
+@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)