introduce class hiearchy, add broadcast to wire successfully

Author: johnzl-777

src/bloqade/squin/qubit.py

@@ -20,25 +20,28 @@
 dialect = ir.Dialect("squin.qubit")
 
 
-@statement(dialect=dialect)
-class New(ir.Statement):
+@statement
+class MultiQubitStatement(ir.Statement):
     traits = frozenset({lowering.FromPythonCall()})
-    n_qubits: ir.SSAValue = info.argument(types.Int)
-    result: ir.ResultValue = info.result(ilist.IListType[QubitType, types.Any])
+    operator: ir.SSAValue = info.argument(OpType)
+    qubits: ir.SSAValue = info.argument(ilist.IListType[QubitType])
 
 
 @statement(dialect=dialect)
-class Apply(ir.Statement):
-    traits = frozenset({lowering.FromPythonCall()})
-    operator: ir.SSAValue = info.argument(OpType)
-    qubits: ir.SSAValue = info.argument(ilist.IListType[QubitType])
+class Apply(MultiQubitStatement):
+    pass
 
 
 @statement(dialect=dialect)
-class Broadcast(ir.Statement):
+class Broadcast(MultiQubitStatement):
+    pass
+
+
+@statement(dialect=dialect)
+class New(ir.Statement):
     traits = frozenset({lowering.FromPythonCall()})
-    operator: ir.SSAValue = info.argument(OpType)
-    qubits: ir.SSAValue = info.argument(ilist.IListType[QubitType])
+    n_qubits: ir.SSAValue = info.argument(types.Int)
+    result: ir.ResultValue = info.result(ilist.IListType[QubitType, types.Any])
 
 
 @statement(dialect=dialect)

src/bloqade/squin/wire.py

@@ -48,10 +48,8 @@ class Unwrap(ir.Statement):
     result: ir.ResultValue = info.result(WireType)
 
 
-# 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, ...)
+@statement
+class MultiWireStatement(ir.Statement):
     traits = frozenset({lowering.FromPythonCall(), ir.Pure()})
     operator: ir.SSAValue = info.argument(OpType)
     inputs: tuple[ir.SSAValue, ...] = info.argument(WireType)
@@ -68,6 +66,20 @@ def __init__(self, operator: ir.SSAValue, *args: ir.SSAValue):
         )  # custom lowering required for wrapper to work here
 
 
+# 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(MultiWireStatement):  # apply(op, w1, w2, ...)
+    def __init__(self, operator: ir.SSAValue, *args: ir.SSAValue):
+        super().__init__(operator, *args)
+
+
+@statement(dialect=dialect)
+class Broadcast(MultiWireStatement):
+    def __init__(self, operator: ir.SSAValue, *args: ir.SSAValue):
+        super().__init__(operator, *args)
+
+
 # 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.