Completed implementation but not quite working

Author: johnzl-777

squin_op_playground.py

@@ -1,8 +1,8 @@
 from kirin import ir, types
 from kirin.passes import Fold
-from kirin.dialects import py, func, ilist
+from kirin.dialects import py, func
 
-from bloqade import qasm2, squin
+from bloqade import squin
 from bloqade.analysis import address
 from bloqade.squin.analysis import shape
 
@@ -11,32 +11,20 @@ def as_int(value: int):
     return py.constant.Constant(value=value)
 
 
-squin_with_qasm_core = squin.groups.wired.add(qasm2.core).add(ilist)
+squin_with_qasm_core = squin.groups.wired
 
 stmts: list[ir.Statement] = [
-    # Create qubit register
-    (n_qubits := as_int(1)),
-    (qreg := qasm2.core.QRegNew(n_qubits=n_qubits.result)),
-    # Get qubits out
-    (idx0 := as_int(0)),
-    (q0 := qasm2.core.QRegGet(reg=qreg.result, idx=idx0.result)),
-    # Unwrap to get wires
-    (w1 := squin.wire.Unwrap(qubit=q0.result)),
-    # Pass wire into operator
-    (op := squin.op.stmts.H()),
-    (v1 := squin.wire.Apply(op.result, w1.result)),
-    # Test Identity
-    (id := squin.op.stmts.Identity(size=1)),
-    (v2 := squin.wire.Apply(id.result, v1.results[0])),
-    # Keep Passing Operators
-    (func.Return(v2.results[0])),
+    (h0 := squin.op.stmts.H()),
+    (h1 := squin.op.stmts.H()),
+    (hh := squin.op.stmts.Kron(lhs=h1.result, rhs=h0.result)),
+    (func.Return(hh.result)),
 ]
 
 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=ilist.IListType),
+    signature=func.Signature(inputs=(), output=squin.op.types.OpType),
     body=ir.Region(blocks=block),
 )
 

src/bloqade/squin/analysis/shape/analysis.py

@@ -23,15 +23,17 @@ def initialize(self):
     # I can get the data I want from the SizedTrait
     # and go from there
     def eval_stmt(self, frame: ForwardFrame, stmt: ir.Statement):
-        if stmt.has_trait(Sized):
-            size = stmt.get_trait(Sized)
-            return (OpShape(size=size.data),)
-        # Handle op.Identity
+        method = self.lookup_registry(frame, stmt)
+        if method is not None:
+            return method(self, frame, stmt)
         elif stmt.has_trait(HasSize):
             # Caution! This can return None
             has_size_inst = stmt.get_trait(HasSize)
             size = has_size_inst.get_size(stmt)
             return (OpShape(size=size),)
+        elif stmt.has_trait(Sized):
+            size = stmt.get_trait(Sized)
+            return (OpShape(size=size.data),)
         else:
             return (NoShape(),)
 

src/bloqade/squin/analysis/shape/impls.py

@@ -1,24 +1,69 @@
-from kirin import interp
+from typing import cast
 
-from bloqade import squin
+from kirin import ir, interp
 
-""" from .lattice import (
-    Shape,
+from bloqade.squin import op
+
+from .lattice import (
     NoShape,
     OpShape,
 )
-
-from .analysis import ShapeAnalysis """
+from .analysis import ShapeAnalysis
 
 
-@squin.op.dialect.register(key="op.shape")
+@op.dialect.register(key="op.shape")
 class SquinOp(interp.MethodTable):
-    pass
 
-    # Should be using the Sized trait
-    # that the statements have
+    @interp.impl(op.stmts.Kron)
+    def kron(self, interp: ShapeAnalysis, frame: interp.Frame, stmt: op.stmts.Kron):
+        lhs = frame.get(stmt.lhs)
+        rhs = frame.get(stmt.rhs)
+        if isinstance(lhs, OpShape) and isinstance(rhs, OpShape):
+            new_size = lhs.size + rhs.size
+            return (OpShape(size=new_size),)
+        else:
+            return (NoShape(),)
+
+    @interp.impl(op.stmts.Mult)
+    def mult(self, interp: ShapeAnalysis, frame: interp.Frame, stmt: op.stmts.Mult):
+        lhs = frame.get(stmt.lhs)
+        rhs = frame.get(stmt.rhs)
+
+        if isinstance(lhs, OpShape) and isinstance(rhs, OpShape):
+            lhs_size = lhs.size
+            rhs_size = rhs.size
+            # Sized trait implicitly enforces that
+            # all operators are square matrices,
+            # not sure if it's worth raising an exception here
+            # or just letting this propagate...
+            if lhs_size != rhs_size:
+                return (NoShape(),)
+            else:
+                return (OpShape(size=lhs_size + rhs_size),)
+        else:
+            return (NoShape(),)
+
+    @interp.impl(op.stmts.Control)
+    def control(
+        self, interp: ShapeAnalysis, frame: interp.Frame, stmt: op.stmts.Control
+    ):
+        op_shape = frame.get(stmt.op)
+
+        if isinstance(op_shape, OpShape):
+            op_size = op_shape.size
+            n_controls_attr = stmt.get_attr_or_prop("n_controls")
+            # raise exception if attribute is NOne
+            n_controls = cast(ir.PyAttr[int], n_controls_attr).data
+            return (OpShape(size=op_size + n_controls),)
+        else:
+            return (NoShape(),)
+
+    @interp.impl(op.stmts.Rot)
+    def rot(self, interp: ShapeAnalysis, frame: interp.Frame, stmt: op.stmts.Rot):
+        op_shape = frame.get(stmt.axis)
+        return op_shape
 
-    # Need to keep in mind that Identity
-    # has a HasSize() trait with "size:int"
-    # as the corresponding attribute to query
-    # @interp.impl(squin.op.stmts.ConstantUnitary)
+    @interp.impl(op.stmts.Scale)
+    def scale(self, interp: ShapeAnalysis, frame: interp.Frame, stmt: op.stmts.Scale):
+        op_shape = frame.get(stmt.op)
+        return op_shape