add unit tests, implement name changes suggested by Roger

Author: johnzl-777

src/bloqade/squin/analysis/nsites/__init__.py

@@ -1,3 +1,8 @@
 # Need this for impl registration to work properly!
 from . import impls as impls
+from .lattice import (
+    NoSites as NoSites,
+    AnySites as AnySites,
+    HasNSites as HasNSites,
+)
 from .analysis import NSitesAnalysis as NSitesAnalysis

src/bloqade/squin/analysis/nsites/analysis.py

@@ -5,15 +5,15 @@
 from kirin.analysis.forward import ForwardFrame
 
 from bloqade.squin.op.types import OpType
-from bloqade.squin.op.traits import Sites, HasNSitesTrait
+from bloqade.squin.op.traits import NSites, HasNSitesTrait
 
-from .lattice import NSites, NoSites, HasNSites
+from .lattice import Sites, NoSites, HasNSites
 
 
-class NSitesAnalysis(Forward[NSites]):
+class NSitesAnalysis(Forward[Sites]):
 
     keys = ["op.nsites"]
-    lattice = NSites
+    lattice = Sites
 
     # Take a page from const prop in Kirin,
     # I can get the data I want from the SizedTrait
@@ -28,16 +28,16 @@ def eval_stmt(self, frame: ForwardFrame, stmt: ir.Statement):
             has_n_sites_trait = stmt.get_trait(HasNSitesTrait)
             sites = has_n_sites_trait.get_sites(stmt)
             return (HasNSites(sites=sites),)
-        elif stmt.has_trait(Sites):
-            sites_trait = stmt.get_trait(Sites)
+        elif stmt.has_trait(NSites):
+            sites_trait = stmt.get_trait(NSites)
             return (HasNSites(sites=sites_trait.data),)
         else:
             return (NoSites(),)
 
     # For when no implementation is found for the statement
     def eval_stmt_fallback(
-        self, frame: ForwardFrame[NSites], stmt: ir.Statement
-    ) -> tuple[NSites, ...]:  # some form of Shape will go back into the frame
+        self, frame: ForwardFrame[Sites], stmt: ir.Statement
+    ) -> tuple[Sites, ...]:  # some form of Shape will go back into the frame
         return tuple(
             (
                 self.lattice.top()
@@ -47,6 +47,6 @@ def eval_stmt_fallback(
             for result in stmt.results
         )
 
-    def run_method(self, method: ir.Method, args: tuple[NSites, ...]):
+    def run_method(self, method: ir.Method, args: tuple[Sites, ...]):
         # NOTE: we do not support dynamic calls here, thus no need to propagate method object
         return self.run_callable(method.code, (self.lattice.bottom(),) + args)

src/bloqade/squin/analysis/nsites/lattice.py

@@ -10,40 +10,40 @@
 
 
 @dataclass
-class NSites(
-    SimpleJoinMixin["NSites"], SimpleMeetMixin["NSites"], BoundedLattice["NSites"]
+class Sites(
+    SimpleJoinMixin["Sites"], SimpleMeetMixin["Sites"], BoundedLattice["Sites"]
 ):
     @classmethod
-    def bottom(cls) -> "NSites":
+    def bottom(cls) -> "Sites":
         return NoSites()
 
     @classmethod
-    def top(cls) -> "NSites":
+    def top(cls) -> "Sites":
         return AnySites()
 
 
 @final
 @dataclass
-class NoSites(NSites, metaclass=SingletonMeta):
+class NoSites(Sites, metaclass=SingletonMeta):
 
-    def is_subseteq(self, other: NSites) -> bool:
+    def is_subseteq(self, other: Sites) -> bool:
         return True
 
 
 @final
 @dataclass
-class AnySites(NSites, metaclass=SingletonMeta):
+class AnySites(Sites, metaclass=SingletonMeta):
 
-    def is_subseteq(self, other: NSites) -> bool:
-        return isinstance(other, NSites)
+    def is_subseteq(self, other: Sites) -> bool:
+        return isinstance(other, Sites)
 
 
 @final
 @dataclass
-class HasNSites(NSites):
+class HasNSites(Sites):
     sites: int
 
-    def is_subseteq(self, other: NSites) -> bool:
+    def is_subseteq(self, other: Sites) -> bool:
         if isinstance(other, HasNSites):
             return self.sites == other.sites
         return False

src/bloqade/squin/op/stmts.py

@@ -2,7 +2,7 @@
 from kirin.decl import info, statement
 
 from .types import OpType
-from .traits import Sites, Unitary, MaybeUnitary, HasNSitesTrait
+from .traits import NSites, Unitary, MaybeUnitary, HasNSitesTrait
 from .complex import Complex
 from ._dialect import dialect
 
@@ -84,14 +84,14 @@ class Identity(CompositeOp):
 
 @statement
 class ConstantOp(PrimitiveOp):
-    traits = frozenset({ir.Pure(), ir.FromPythonCall(), ir.ConstantLike(), Sites(1)})
+    traits = frozenset({ir.Pure(), ir.FromPythonCall(), ir.ConstantLike(), NSites(1)})
     result: ir.ResultValue = info.result(OpType)
 
 
 @statement
 class ConstantUnitary(ConstantOp):
     traits = frozenset(
-        {ir.Pure(), ir.FromPythonCall(), ir.ConstantLike(), Unitary(), Sites(1)}
+        {ir.Pure(), ir.FromPythonCall(), ir.ConstantLike(), Unitary(), NSites(1)}
     )
 
 
@@ -105,7 +105,7 @@ class PhaseOp(PrimitiveOp):
     $$
     """
 
-    traits = frozenset({ir.Pure(), ir.FromPythonCall(), Unitary(), Sites(1)})
+    traits = frozenset({ir.Pure(), ir.FromPythonCall(), Unitary(), NSites(1)})
     theta: ir.SSAValue = info.argument(types.Float)
     result: ir.ResultValue = info.result(OpType)
 
@@ -120,7 +120,7 @@ class ShiftOp(PrimitiveOp):
     $$
     """
 
-    traits = frozenset({ir.Pure(), ir.FromPythonCall(), Unitary(), Sites(1)})
+    traits = frozenset({ir.Pure(), ir.FromPythonCall(), Unitary(), NSites(1)})
     theta: ir.SSAValue = info.argument(types.Float)
     result: ir.ResultValue = info.result(OpType)
 

src/bloqade/squin/op/traits.py

@@ -5,7 +5,7 @@
 
 
 @dataclass(frozen=True)
-class Sites(ir.StmtTrait):
+class NSites(ir.StmtTrait):
     data: int
 
 

test/squin/analysis/test_nsites_analysis.py

@@ -0,0 +1,225 @@
+from kirin import ir, types
+from kirin.passes import Fold
+from kirin.dialects import py, func
+
+from bloqade import squin
+from bloqade.squin.analysis import nsites
+
+
+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_primitive_ops():
+    pass
+
+
+# Kron, Mult, Control, Rot, and Scale all have methods defined for handling them in impls,
+# The following should ensure the code paths are properly exercised
+
+
+def test_control():
+    # Control doesn't have an impl but it is handled in the eval_stmt of the interpreter
+    # because it has a HasNSitesTrait future statements might have
+
+    stmts: list[ir.Statement] = [
+        (h0 := squin.op.stmts.H()),
+        (controlled_h := squin.op.stmts.Control(op=h0.result, n_controls=1)),
+        (func.Return(controlled_h.result)),
+    ]
+
+    constructed_method = gen_func_from_stmts(stmts)
+
+    nsites_frame, _ = nsites.NSitesAnalysis(constructed_method.dialects).run_analysis(
+        constructed_method, no_raise=False
+    )
+
+    has_n_sites = []
+    for nsites_type in nsites_frame.entries.values():
+        if isinstance(nsites_type, nsites.HasNSites):
+            has_n_sites.append(nsites_type)
+
+    assert len(has_n_sites) == 2
+    assert has_n_sites[0].sites == 1
+    assert has_n_sites[1].sites == 2
+
+
+def test_kron():
+
+    stmts: list[ir.Statement] = [
+        (h0 := squin.op.stmts.H()),
+        (h1 := squin.op.stmts.H()),
+        (hh := squin.op.stmts.Kron(h0.result, h1.result)),
+        (func.Return(hh.result)),
+    ]
+
+    constructed_method = gen_func_from_stmts(stmts)
+
+    nsites_frame, _ = nsites.NSitesAnalysis(constructed_method.dialects).run_analysis(
+        constructed_method, no_raise=False
+    )
+
+    has_n_sites = []
+    for nsites_type in nsites_frame.entries.values():
+        if isinstance(nsites_type, nsites.HasNSites):
+            has_n_sites.append(nsites_type)
+
+    assert len(has_n_sites) == 3
+    assert has_n_sites[0].sites == 1
+    assert has_n_sites[1].sites == 1
+    assert has_n_sites[2].sites == 3
+
+
+def test_mult_square_same_sites():
+    # Ensure that two operators of the same size produce
+    # a valid operator as their result
+
+    stmts: list[ir.Statement] = [
+        (h0 := squin.op.stmts.H()),
+        (h1 := squin.op.stmts.H()),
+        (h2 := squin.op.stmts.Mult(h0.result, h1.result)),
+        (func.Return(h2.result)),
+    ]
+
+    constructed_method = gen_func_from_stmts(stmts)
+
+    nsites_frame, _ = nsites.NSitesAnalysis(constructed_method.dialects).run_analysis(
+        constructed_method, no_raise=False
+    )
+
+    has_n_sites = []
+    for nsites_type in nsites_frame.entries.values():
+        if isinstance(nsites_type, nsites.HasNSites):
+            has_n_sites.append(nsites_type)
+
+    # should be three HasNSites types
+    assert len(has_n_sites) == 3
+    # the first 2 HasNSites will have 1 site but
+    # the Kron-produced operator should have 2 sites
+    assert has_n_sites[0].sites == 1
+    assert has_n_sites[1].sites == 1
+    assert has_n_sites[2].sites == 2
+
+
+def test_mult_square_different_sites():
+    # Ensure that two operators of different sizes produce
+    # NoSites as a type. Note that a better solution would be
+    # to implement a special error type in the type lattice
+    # but this would introduce some complexity later on
+
+    stmts: list[ir.Statement] = [
+        (h0 := squin.op.stmts.H()),
+        (h1 := squin.op.stmts.H()),
+        # Kron to make nsites = 2 operator
+        (hh := squin.op.stmts.Kron(h0.result, h1.result)),
+        # apply Mult on HasNSites(2) and HasNSites(1)
+        (invalid_op := squin.op.stmts.Mult(hh.result, h1.result)),
+        (func.Return(invalid_op.result)),
+    ]
+
+    constructed_method = gen_func_from_stmts(stmts)
+
+    nsites_frame, _ = nsites.NSitesAnalysis(constructed_method.dialects).run_analysis(
+        constructed_method, no_raise=False
+    )
+
+    nsites_types = list(nsites_frame.entries.values())
+
+    has_n_sites = []
+    no_sites = []
+    for nsite_type in nsites_types:
+        if isinstance(nsite_type, nsites.HasNSites):
+            has_n_sites.append(nsite_type)
+        elif isinstance(nsite_type, nsites.NoSites):
+            no_sites.append(nsite_type)
+
+    assert len(has_n_sites) == 3
+    # HasNSites(1) for Hadamards, 2 for Kron result
+    assert has_n_sites[0].sites == 1
+    assert has_n_sites[1].sites == 1
+    assert has_n_sites[2].sites == 2
+    # One from function itself, another from invalid mult
+    assert len(no_sites) == 2
+
+
+def test_rot():
+
+    stmts: list[ir.Statement] = [
+        (h0 := squin.op.stmts.H()),
+        (angle := as_float(0.2)),
+        (rot_h := squin.op.stmts.Rot(axis=h0.result, angle=angle.result)),
+        (func.Return(rot_h.result)),
+    ]
+
+    constructed_method = gen_func_from_stmts(stmts)
+
+    nsites_frame, _ = nsites.NSitesAnalysis(constructed_method.dialects).run_analysis(
+        constructed_method, no_raise=False
+    )
+
+    has_n_sites = []
+    for nsites_type in nsites_frame.entries.values():
+        if isinstance(nsites_type, nsites.HasNSites):
+            has_n_sites.append(nsites_type)
+
+    assert len(has_n_sites) == 2
+    # Rot should just propagate whatever Sites type is there
+    assert has_n_sites[0].sites == 1
+    assert has_n_sites[1].sites == 1
+
+
+def test_scale():
+
+    stmts: list[ir.Statement] = [
+        (h0 := squin.op.stmts.H()),
+        (factor := as_float(0.2)),
+        (rot_h := squin.op.stmts.Scale(op=h0.result, factor=factor.result)),
+        (func.Return(rot_h.result)),
+    ]
+
+    constructed_method = gen_func_from_stmts(stmts)
+
+    nsites_frame, _ = nsites.NSitesAnalysis(constructed_method.dialects).run_analysis(
+        constructed_method, no_raise=False
+    )
+
+    has_n_sites = []
+    for nsites_type in nsites_frame.entries.values():
+        if isinstance(nsites_type, nsites.HasNSites):
+            has_n_sites.append(nsites_type)
+
+    assert len(has_n_sites) == 2
+    # Rot should just propagate whatever Sites type is there
+    assert has_n_sites[0].sites == 1
+    assert has_n_sites[1].sites == 1