fix more scf bugs for bloqade (#243)

Author: Roger-luo

src/kirin/dialects/scf/__init__.py

@@ -11,6 +11,7 @@
 """
 
 from . import (
+    absint as absint,
     interp as interp,
     unroll as unroll,
     lowering as lowering,

src/kirin/dialects/scf/absint.py

@@ -0,0 +1,62 @@
+from kirin import ir, types, interp
+from kirin.analysis import const
+from kirin.dialects import func
+
+from .stmts import Yield, IfElse
+from ._dialect import dialect
+
+
+@dialect.register(key="absint")
+class Methods(interp.MethodTable):
+
+    @interp.impl(Yield)
+    def yield_stmt(
+        self,
+        interp_: interp.AbstractInterpreter,
+        frame: interp.AbstractFrame,
+        stmt: Yield,
+    ):
+        return interp.YieldValue(frame.get_values(stmt.values))
+
+    @interp.impl(IfElse)
+    def if_else(
+        self,
+        interp_: interp.AbstractInterpreter,
+        frame: interp.AbstractFrame,
+        stmt: IfElse,
+    ):
+        if isinstance(hint := stmt.cond.hints.get("const"), const.Value):
+            if hint.data:
+                return self._infer_if_else_cond(interp_, frame, stmt, stmt.then_body)
+            else:
+                return self._infer_if_else_cond(interp_, frame, stmt, stmt.else_body)
+        then_results = self._infer_if_else_cond(interp_, frame, stmt, stmt.then_body)
+        else_results = self._infer_if_else_cond(interp_, frame, stmt, stmt.else_body)
+
+        match (then_results, else_results):
+            case (interp.ReturnValue(then_value), interp.ReturnValue(else_value)):
+                return interp.ReturnValue(then_value.join(else_value))
+            case (interp.ReturnValue(then_value), _):
+                return then_results
+            case (_, interp.ReturnValue(else_value)):
+                return else_results
+            case _:
+                return interp_.join_results(then_results, else_results)
+
+    def _infer_if_else_cond(
+        self,
+        interp_: interp.AbstractInterpreter,
+        frame: interp.AbstractFrame,
+        stmt: IfElse,
+        body: ir.Region,
+    ):
+        body_block = body.blocks[0]
+        body_term = body_block.last_stmt
+        if isinstance(body_term, func.Return):
+            frame.worklist.append(interp.Successor(body_block, types.Bool))
+            return
+
+        with interp_.state.new_frame(interp_.new_frame(stmt)) as body_frame:
+            body_frame.entries.update(frame.entries)
+            body_frame.set(body_block.args[0], types.Bool)
+            return interp_.run_ssacfg_region(body_frame, stmt.then_body)

src/kirin/dialects/scf/constprop.py

@@ -2,6 +2,7 @@
 
 from kirin import ir, interp
 from kirin.analysis import const
+from kirin.dialects import func
 
 from .stmts import For, Yield, IfElse
 from ._dialect import dialect
@@ -48,16 +49,23 @@ def if_else(
             else_frame, else_results = self._prop_const_cond_ifelse(
                 interp_, frame, stmt, const.Value(False), stmt.else_body
             )
-            ret = interp_.join_results(then_results, else_results)
-
-            if not then_frame.frame_is_not_pure or not else_frame.frame_is_not_pure:
-                frame.should_be_pure.add(stmt)
-
             # NOTE: then_frame and else_frame do not change
             # parent frame variables value except cond
             frame.entries.update(then_frame.entries)
             frame.entries.update(else_frame.entries)
-            frame.set(stmt.cond, cond)
+            # TODO: pick the non-return value
+            if isinstance(then_results, interp.ReturnValue) and isinstance(
+                else_results, interp.ReturnValue
+            ):
+                return interp.ReturnValue(then_results.value.join(else_results.value))
+            elif isinstance(then_results, interp.ReturnValue):
+                ret = else_results
+            elif isinstance(else_results, interp.ReturnValue):
+                ret = then_results
+            else:
+                if not then_frame.frame_is_not_pure or not else_frame.frame_is_not_pure:
+                    frame.should_be_pure.add(stmt)
+                ret = interp_.join_results(then_results, else_results)
         return ret
 
     def _prop_const_cond_ifelse(
@@ -73,7 +81,9 @@ def _prop_const_cond_ifelse(
             body_frame.set(body.blocks[0].args[0], cond)
             results = interp_.run_ssacfg_region(body_frame, body)
 
-        if not body_frame.frame_is_not_pure:
+        if not body_frame.frame_is_not_pure and not isinstance(
+            body.blocks[0].last_stmt, func.Return
+        ):
             frame.should_be_pure.add(stmt)
         return body_frame, results
 

src/kirin/dialects/scf/typeinfer.py

@@ -1,26 +1,18 @@
 from kirin import ir, types, interp
-from kirin.analysis import ForwardFrame, TypeInference, const
+from kirin.analysis import ForwardFrame, TypeInference
 from kirin.dialects import func
 from kirin.dialects.eltype import ElType
 
-from .stmts import For, Yield, IfElse
+from . import absint
+from .stmts import For, IfElse
 from ._dialect import dialect
 
 
 @dialect.register(key="typeinfer")
-class TypeInfer(interp.MethodTable):
-
-    @interp.impl(Yield)
-    def yield_stmt(
-        self,
-        interp_: TypeInference,
-        frame: ForwardFrame[types.TypeAttribute],
-        stmt: Yield,
-    ):
-        return interp.YieldValue(frame.get_values(stmt.values))
+class TypeInfer(absint.Methods):
 
     @interp.impl(IfElse)
-    def if_else(
+    def if_else_(
         self,
         interp_: TypeInference,
         frame: ForwardFrame[types.TypeAttribute],
@@ -29,32 +21,7 @@ def if_else(
         frame.set(
             stmt.cond, frame.get(stmt.cond).meet(types.Bool)
         )  # set cond backwards
-        if isinstance(hint := stmt.cond.hints.get("const"), const.Value):
-            if hint.data:
-                return self._infer_if_else_cond(interp_, frame, stmt, stmt.then_body)
-            else:
-                return self._infer_if_else_cond(interp_, frame, stmt, stmt.else_body)
-        then_results = self._infer_if_else_cond(interp_, frame, stmt, stmt.then_body)
-        else_results = self._infer_if_else_cond(interp_, frame, stmt, stmt.else_body)
-        return interp_.join_results(then_results, else_results)
-
-    def _infer_if_else_cond(
-        self,
-        interp_: TypeInference,
-        frame: ForwardFrame[types.TypeAttribute],
-        stmt: IfElse,
-        body: ir.Region,
-    ):
-        body_block = body.blocks[0]
-        body_term = body_block.last_stmt
-        if isinstance(body_term, func.Return):  # TODO: use types.Literal?
-            frame.worklist.append(interp.Successor(body_block, types.Bool))
-            return
-
-        with interp_.state.new_frame(interp_.new_frame(stmt)) as body_frame:
-            body_frame.entries.update(frame.entries)
-            body_frame.set(body_block.args[0], types.Bool)
-            return interp_.run_ssacfg_region(body_frame, stmt.then_body)
+        return super().if_else(self, interp_, frame, stmt)
 
     @interp.impl(For)
     def for_loop(

src/kirin/interp/impl.py

@@ -70,6 +70,9 @@ def __repr__(self):
         else:
             return f"interp {self.parent.name}"
 
+    def __call__(self, *args, **kwargs):
+        return self.impl(*args, **kwargs)
+
 
 @dataclass
 class AttributeImplDef(Def[type[Attribute], "AttributeFunction"]):