make py module independent

Author: Roger-luo

src/kirin/analysis/forward.py

@@ -27,6 +27,7 @@ def __init__(
         dialects: ir.DialectGroup | Iterable[ir.Dialect],
         *,
         fuel: int | None = None,
+        save_all_ssa: bool = False,
         max_depth: int = 128,
         max_python_recursion_depth: int = 8192,
     ):
@@ -36,7 +37,8 @@ def __init__(
             max_depth=max_depth,
             max_python_recursion_depth=max_python_recursion_depth,
         )
-        self.results = {}
+        self.save_all_ssa = save_all_ssa
+        self.results: dict[ir.SSAValue, LatticeElemType] = {}
 
     def set_values(
         self,
@@ -55,7 +57,10 @@ def finalize(
         frame: ForwardFrame[LatticeElemType, ExtraType],
         results: MethodResult[LatticeElemType],
     ) -> MethodResult[LatticeElemType]:
-        self.results = frame.entries
+        if self.save_all_ssa:
+            self.results.update(frame.entries)
+        else:
+            self.results = frame.entries
         return results
 
     def new_frame(self, code: ir.Statement) -> ForwardFrame[LatticeElemType, ExtraType]:

src/kirin/decl/emit/from_python_call.py

@@ -6,7 +6,6 @@
 from kirin.lowering import Result, LoweringState
 from kirin.decl.base import BaseModifier, StatementOptions
 from kirin.exceptions import DialectLoweringError
-from kirin.dialects.py import data
 
 from ._create_fn import create_fn
 from ._set_new_attribute import set_new_attribute
@@ -29,6 +28,8 @@ class EmitFromPythonCall(BaseModifier):
     _AST_CONSTANT = "_kirin_ast_Constant"
 
     def __init__(self, cls: type, **kwargs: Unpack[StatementOptions]) -> None:
+        from kirin.dialects.py import data
+
         super().__init__(cls, **kwargs)
         self.globals[self._KIRIN_RESULT] = Result
         self.globals[self._KIRIN_PYATTR] = data.PyAttr

src/kirin/decl/emit/init.py

@@ -6,7 +6,6 @@
 from kirin import ir
 from kirin.decl import info
 from kirin.decl.base import BaseModifier, StatementOptions
-from kirin.dialects.py.data import PyAttr
 
 from ._create_fn import create_fn
 from ._set_new_attribute import set_new_attribute
@@ -30,6 +29,8 @@ class EmitInit(BaseModifier):
 
     def __init__(self, cls: type, **kwargs: Unpack[StatementOptions]) -> None:
         super().__init__(cls, **kwargs)
+        from kirin.dialects.py.data import PyAttr
+
         self._init_params: list[str] = []
         self._init_body: list[str] = []
         self._init_locals: dict[str, Any] = {}

src/kirin/decl/emit/property.py

@@ -1,6 +1,5 @@
 from kirin import ir
 from kirin.decl import info
-from kirin.dialects.py import data
 from kirin.decl.emit.init import BaseModifier
 
 from ._create_fn import create_fn
@@ -112,6 +111,8 @@ def _emit_result_property(self, index: int, f: info.ResultField):
         return getter, setter
 
     def _emit_attribute_property(self, f: info.AttributeField):
+        from kirin.dialects.py import data
+
         storage = "properties" if f.property else "attributes"
         attr = f"{self._self_name}.{storage}['{f.name}']"
         getter = create_fn(

src/kirin/dialects/py/__init__.py

@@ -1,4 +1,5 @@
-# from . import data as data, stmts as stmts
+from . import data as data, stmts as stmts
+
 # from kirin.dialects.pytypes import types as types
 # from kirin.dialects.pytypes.data import PyAttr as PyAttr
 # from kirin.dialects.pytypes.types.base import (