refactor analysis with SSAValue hints (#234)

## Analysis now returns frame instead of `dict`

Previously we record the results into a global field `analysis.results`
by merging the `frame.entries` into this dictionary. This works for
simple cases, but in the case of constant prop, the purity flag is per
frame (e.g when you go through one of the successor branch/region).

As a result, when a function call has no return value, or a scf.For does
not have yield (thus no results). We will not mark it pure.

More generally, when there are extra results per function, we won't be
able to return that results. So now we return the frame we use to run
the analysis within, which contains everything one defines to collect
for the analysis.

## Fixing DCE on SCF

with the above change, we are able to fix DCE on SCF and DCE now only
checks `Pure` or `MaybePure`.

## Replacing `Hint` with `SSAValue.hints`

A new field `hints` is added to `SSAValue`. This is because I realize
there could be different kind of hints from analysis to carry over
within the IR. This also simplifies the type system making it less
complicated.

---

well this turns out to be quite breaking... But it simplifies things a
lot.

This also fixes a few bugs related to purity analysis - the purity
should be per frame/statement instead of per SSAValue (well... I always
felt this was strange, now it is correct). Constant prop is no longer
convoluted with purity analysis (still slightly but a lot easier to
decouple in the future with joint analysis).

As a result loops on pure functions can fold correctly now.

Author: Roger-luo

src/kirin/__init__.py

@@ -1,6 +1,7 @@
 # re-exports the public API of the kirin package
 from kirin import ir
-from kirin.ir import types as types
 from kirin.decl import info, statement
 
-__all__ = ["ir", "statement", "info"]
+from . import types as types
+
+__all__ = ["ir", "types", "statement", "info"]

src/kirin/analysis/const/__init__.py

@@ -8,17 +8,13 @@
 the IR.
 """
 
-from .prop import Propagate as Propagate, ExtraFrameInfo as ExtraFrameInfo
+from .prop import Frame as Frame, Propagate as Propagate
 from .lattice import (
-    Pure as Pure,
     Value as Value,
     Bottom as Bottom,
     Result as Result,
-    NotPure as NotPure,
     Unknown as Unknown,
-    JointResult as JointResult,
     PartialConst as PartialConst,
     PartialTuple as PartialTuple,
-    PurityBottom as PurityBottom,
     PartialLambda as PartialLambda,
 )

src/kirin/analysis/const/lattice.py

@@ -2,28 +2,30 @@
 """
 
 from typing import Any, final
-from dataclasses import field, dataclass
+from dataclasses import dataclass
 
 from kirin import ir
 from kirin.lattice import (
-    LatticeMeta,
-    SingletonMeta,
     BoundedLattice,
     IsSubsetEqMixin,
     SimpleJoinMixin,
     SimpleMeetMixin,
 )
+from kirin.ir.attrs.abc import LatticeAttributeMeta, SingletonLatticeAttributeMeta
+from kirin.print.printer import Printer
 
 from ._visitor import _ElemVisitor
 
 
 @dataclass
 class Result(
+    ir.Attribute,
     IsSubsetEqMixin["Result"],
     SimpleJoinMixin["Result"],
     SimpleMeetMixin["Result"],
     BoundedLattice["Result"],
     _ElemVisitor,
+    metaclass=LatticeAttributeMeta,
 ):
     """Base class for constant analysis results."""
 
@@ -35,29 +37,38 @@ def top(cls) -> "Result":
     def bottom(cls) -> "Result":
         return Bottom()
 
+    def print_impl(self, printer: Printer) -> None:
+        printer.plain_print(repr(self))
+
 
 @final
 @dataclass
-class Unknown(Result, metaclass=SingletonMeta):
+class Unknown(Result, metaclass=SingletonLatticeAttributeMeta):
     """Unknown constant value. This is the top element of the lattice."""
 
     def is_subseteq(self, other: Result) -> bool:
         return isinstance(other, Unknown)
 
+    def __hash__(self) -> int:
+        return id(self)
+
 
 @final
 @dataclass
-class Bottom(Result, metaclass=SingletonMeta):
+class Bottom(Result, metaclass=SingletonLatticeAttributeMeta):
     """Bottom element of the lattice."""
 
     def is_subseteq(self, other: Result) -> bool:
         return True
 
+    def __hash__(self) -> int:
+        return id(self)
+
 
 @final
 @dataclass
 class Value(Result):
-    """Constant value. Wraps any Python value."""
+    """Constant value. Wraps any hashable Python value."""
 
     data: Any
 
@@ -69,6 +80,12 @@ def is_equal(self, other: Result) -> bool:
             return self.data == other.data
         return False
 
+    def __hash__(self) -> int:
+        # NOTE: we use id here because the data
+        # may not be hashable. This is fine because
+        # the data is guaranteed to be unique.
+        return id(self)
+
 
 @dataclass
 class PartialConst(Result):
@@ -78,7 +95,7 @@ class PartialConst(Result):
 
 
 @final
-class PartialTupleMeta(LatticeMeta):
+class PartialTupleMeta(LatticeAttributeMeta):
     """Metaclass for PartialTuple.
 
     This metaclass canonicalizes PartialTuple instances with all Value elements
@@ -139,6 +156,9 @@ def is_subseteq_Value(self, other: Value) -> bool:
             return all(x.is_subseteq(Value(y)) for x, y in zip(self.data, other.data))
         return False
 
+    def __hash__(self) -> int:
+        return hash(self.data)
+
 
 @final
 @dataclass
@@ -152,6 +172,9 @@ class PartialLambda(PartialConst):
     code: ir.Statement
     captured: tuple[Result, ...]
 
+    def __hash__(self) -> int:
+        return hash((self.argnames, self.code, self.captured))
+
     def is_subseteq_PartialLambda(self, other: "PartialLambda") -> bool:
         if self.code is not other.code:
             return False
@@ -194,82 +217,3 @@ def meet(self, other: Result) -> Result:
             self.code,
             tuple(x.meet(y) for x, y in zip(self.captured, other.captured)),
         )
-
-
-@dataclass(frozen=True)
-class Purity(
-    SimpleJoinMixin["Purity"], SimpleMeetMixin["Purity"], BoundedLattice["Purity"]
-):
-    """Base class for purity lattice."""
-
-    @classmethod
-    def bottom(cls) -> "Purity":
-        return PurityBottom()
-
-    @classmethod
-    def top(cls) -> "Purity":
-        return NotPure()
-
-
-@dataclass(frozen=True)
-class Pure(Purity, metaclass=SingletonMeta):
-    """The result is from a pure function."""
-
-    def is_subseteq(self, other: Purity) -> bool:
-        return isinstance(other, (NotPure, Pure))
-
-
-@dataclass(frozen=True)
-class NotPure(Purity, metaclass=SingletonMeta):
-    """The result is from an impure function."""
-
-    def is_subseteq(self, other: Purity) -> bool:
-        return isinstance(other, NotPure)
-
-
-@dataclass(frozen=True)
-class PurityBottom(Purity, metaclass=SingletonMeta):
-    """The bottom element of the purity lattice."""
-
-    def is_subseteq(self, other: Purity) -> bool:
-        return True
-
-
-@dataclass
-class JointResult(BoundedLattice["JointResult"]):
-    """Joint result of constant value and purity.
-
-    This lattice is used to join the constant value and purity of a function
-    during constant propagation analysis. This allows the analysis to track
-    both the constant value and the purity of the function, so that the analysis
-    can propagate constant values through function calls even if the function
-    is only partially pure.
-    """
-
-    const: Result
-    """The constant value of the result."""
-    purity: Purity = field(default_factory=Purity.top)
-    """The purity of statement that produces the result."""
-
-    @classmethod
-    def from_const(cls, value: Any) -> "JointResult":
-        return cls(Value(value), Purity.top())
-
-    @classmethod
-    def top(cls) -> "JointResult":
-        return cls(Result.top(), Purity.top())
-
-    @classmethod
-    def bottom(cls) -> "JointResult":
-        return cls(Result.bottom(), Purity.bottom())
-
-    def is_subseteq(self, other: "JointResult") -> bool:
-        return self.const.is_subseteq(other.const) and self.purity.is_subseteq(
-            other.purity
-        )
-
-    def join(self, other: "JointResult") -> "JointResult":
-        return JointResult(self.const.join(other.const), self.purity.join(other.purity))
-
-    def meet(self, other: "JointResult") -> "JointResult":
-        return JointResult(self.const.meet(other.const), self.purity.join(other.purity))

src/kirin/analysis/const/prop.py

@@ -1,22 +1,25 @@
 from dataclasses import field, dataclass
 
-from kirin import ir, interp
+from kirin import ir, types, interp
 from kirin.analysis.forward import ForwardExtra, ForwardFrame
 
-from .lattice import Pure, Value, NotPure, Unknown, JointResult
+from .lattice import Value, Result, Unknown
 
 
 @dataclass
-class ExtraFrameInfo:
+class Frame(ForwardFrame[Result]):
+    should_be_pure: set[ir.Statement] = field(default_factory=set)
+    """If any ir.MaybePure is actually pure."""
     frame_is_not_pure: bool = False
+    """If we hit any non-pure statement."""
 
 
 @dataclass
-class Propagate(ForwardExtra[JointResult, ExtraFrameInfo]):
+class Propagate(ForwardExtra[Frame, Result]):
     """Forward dataflow analysis for constant propagation.
 
     This analysis is a forward dataflow analysis that propagates constant values
-    through the program. It uses the `JointResult` lattice to track the constant
+    through the program. It uses the `Result` lattice to track the constant
     values and purity of the values.
 
     The analysis is implemented as a forward dataflow analysis, where the
@@ -30,7 +33,7 @@ class Propagate(ForwardExtra[JointResult, ExtraFrameInfo]):
     """
 
     keys = ["constprop"]
-    lattice = JointResult
+    lattice = Result
 
     _interp: interp.Interpreter = field(init=False)
 
@@ -49,93 +52,65 @@ def initialize(self):
         self._interp.initialize()
         return self
 
+    def new_frame(self, code: ir.Statement) -> Frame:
+        return Frame.from_func_like(code)
+
     def _try_eval_const_pure(
         self,
-        frame: ForwardFrame[JointResult, ExtraFrameInfo],
+        frame: Frame,
         stmt: ir.Statement,
         values: tuple[Value, ...],
-    ) -> interp.StatementResult[JointResult]:
+    ) -> interp.StatementResult[Result]:
         try:
             _frame = self._interp.new_frame(frame.code)
             _frame.set_values(stmt.args, tuple(x.data for x in values))
             value = self._interp.eval_stmt(_frame, stmt)
             match value:
                 case tuple():
-                    return tuple(JointResult(Value(each), Pure()) for each in value)
+                    return tuple(Value(each) for each in value)
                 case interp.ReturnValue(ret):
-                    return interp.ReturnValue(JointResult(Value(ret), Pure()))
+                    return interp.ReturnValue(Value(ret))
                 case interp.YieldValue(yields):
-                    return interp.YieldValue(
-                        tuple(JointResult(Value(each), Pure()) for each in yields)
-                    )
+                    return interp.YieldValue(tuple(Value(each) for each in yields))
                 case interp.Successor(block, args):
                     return interp.Successor(
                         block,
-                        *tuple(JointResult(Value(each), Pure()) for each in args),
+                        *tuple(Value(each) for each in args),
                     )
         except interp.InterpreterError:
             pass
         return (self.void,)
 
     def eval_stmt(
-        self, frame: ForwardFrame[JointResult, ExtraFrameInfo], stmt: ir.Statement
-    ) -> interp.StatementResult[JointResult]:
+        self, frame: Frame, stmt: ir.Statement
+    ) -> interp.StatementResult[Result]:
         if stmt.has_trait(ir.ConstantLike):
             return self._try_eval_const_pure(frame, stmt, ())
         elif stmt.has_trait(ir.Pure):
-            values = tuple(x.const for x in frame.get_values(stmt.args))
-            if ir.types.is_tuple_of(values, Value):
+            values = frame.get_values(stmt.args)
+            if types.is_tuple_of(values, Value):
                 return self._try_eval_const_pure(frame, stmt, values)
 
         method = self.lookup_registry(frame, stmt)
-        if method is not None:
-            ret = method(self, frame, stmt)
-            self._set_frame_not_pure(ret)
+        if method is None:
+            if stmt.has_trait(ir.Pure):
+                return (Unknown(),)  # no implementation but pure
+            # not pure, and no implementation, let's say it's not pure
+            frame.frame_is_not_pure = True
+            return (Unknown(),)
+
+        ret = method(self, frame, stmt)
+        if stmt.has_trait(ir.IsTerminator):
             return ret
-        elif stmt.has_trait(ir.Pure):
-            # fallback to top for other statements
-            return (JointResult(Unknown(), Pure()),)
-        else:
-            if frame.extra is None:
-                frame.extra = ExtraFrameInfo(True)
-            return (JointResult(Unknown(), NotPure()),)
-
-    def _set_frame_not_pure(self, result: interp.StatementResult[JointResult]):
-        frame = self.state.current_frame()
-        if isinstance(result, tuple) and all(x.purity is Pure() for x in result):
-            return
-
-        if isinstance(result, interp.ReturnValue) and isinstance(
-            result.value.purity, Pure
-        ):
-            return
-
-        if isinstance(result, interp.YieldValue) and all(
-            isinstance(x.purity, Pure) for x in result
-        ):
-            return
-
-        if isinstance(result, interp.Successor) and all(
-            x.purity is Pure() for x in result.block_args
-        ):
-            return
-
-        if frame.extra is None:
-            frame.extra = ExtraFrameInfo(True)
+        elif not stmt.has_trait(ir.MaybePure):  # cannot be pure at all
+            frame.frame_is_not_pure = True
+        elif (
+            stmt not in frame.should_be_pure
+        ):  # implementation cannot decide if it's pure
+            frame.frame_is_not_pure = True
+        return ret
 
     def run_method(
-        self, method: ir.Method, args: tuple[JointResult, ...]
-    ) -> JointResult:
-        return self.run_callable(
-            method.code, (JointResult(Value(method), NotPure()),) + args
-        )
-
-    def finalize(
-        self,
-        frame: ForwardFrame[JointResult, ExtraFrameInfo],
-        results: JointResult,
-    ) -> JointResult:
-        results = super().finalize(frame, results)
-        if frame.extra is not None and frame.extra.frame_is_not_pure:
-            return JointResult(results.const, NotPure())
-        return results
+        self, method: ir.Method, args: tuple[Result, ...]
+    ) -> tuple[Frame, Result]:
+        return self.run_callable(method.code, (Value(method),) + args)