Inline len for IList (#316)

Some kinds of `IList` can't be folded but the length can be determined
at compile time. This PR adds a rewrite that inlines the `len` if the
length can be determined by type inference.

This is preventing some kinds of bloqade programs from being folded to
QASM2. The list of qubits can't be folded to a constant value because
there is no concrete interpreter for `QRegGet`.


```python


@qasm2.extended
def _log_ghz(i_layer: int, qubits: ilist.IList[qasm2.Qubit, Any]):
    stride = len(qubits) // (2**i_layer)

    if stride == 0:
        return
    
    offset = stride // 2
    for j in ilist.range(0, len(qubits), stride):
        qasm2.cx(ctrl=qubits[j], qarg=qubits[j + offset])

    _log_ghz(i_layer + 1, qubits)
    

@qasm2.extended
def main():
    q = qasm2.qreg(6)
    _log_ghz(0,[q[0], q[2], q[4], q[1], q[3], q[5]])
    

fold = QASM2Fold(qasm2.extended)
fold(main)
main.print()
fold(main)
```

The second fold fails because of this.

---------

Co-authored-by: Xiu-zhe (Roger) Luo <[email protected]>

Author: weinbe58

src/kirin/dialects/ilist/rewrite/__init__.py

@@ -1,3 +1,4 @@
 from .list import List2IList as List2IList
 from .const import ConstList2IList as ConstList2IList
 from .unroll import Unroll as Unroll
+from .hint_len import HintLen as HintLen

src/kirin/dialects/ilist/rewrite/hint_len.py

@@ -0,0 +1,34 @@
+from kirin import ir, types
+from kirin.analysis import const
+from kirin.dialects import py
+from kirin.rewrite.abc import RewriteRule, RewriteResult
+from kirin.dialects.ilist.stmts import IListType
+
+
+class HintLen(RewriteRule):
+
+    def _get_collection_len(self, collection: ir.SSAValue):
+        coll_type = collection.type
+
+        if not isinstance(coll_type, types.Generic):
+            return None
+
+        if (
+            coll_type.is_subseteq(IListType)
+            and isinstance(coll_type.vars[1], types.Literal)
+            and isinstance(coll_type.vars[1].data, int)
+        ):
+            return coll_type.vars[1].data
+        else:
+            return None
+
+    def rewrite_Statement(self, node: ir.Statement) -> RewriteResult:
+        if not isinstance(node, py.Len):
+            return RewriteResult()
+
+        if (coll_len := self._get_collection_len(node.value)) is None:
+            return RewriteResult()
+
+        node.result.hints["const"] = const.Value(coll_len)
+
+        return RewriteResult(has_done_something=True)

src/kirin/passes/__init__.py

@@ -1,3 +1,6 @@
 from kirin.passes.abc import Pass as Pass
 from kirin.passes.fold import Fold as Fold
 from kirin.passes.typeinfer import TypeInfer as TypeInfer
+
+from .default import Default as Default
+from .hint_const import HintConst as HintConst

src/kirin/passes/abc.py

@@ -24,7 +24,7 @@ class Pass(ABC):
 
     name: ClassVar[str]
     dialects: DialectGroup
-    no_raise: bool = field(default=False, kw_only=True)
+    no_raise: bool = field(default=True, kw_only=True)
 
     def __call__(self, mt: Method) -> RewriteResult:
         result = self.unsafe_run(mt)
@@ -38,7 +38,7 @@ def fixpoint(self, mt: Method, max_iter: int = 32) -> RewriteResult:
             result = result_.join(result)
             if not result.has_done_something:
                 break
-        mt.code.verify()
+        mt.verify()
         return result
 
     @abstractmethod

src/kirin/passes/aggressive/fold.py

@@ -6,30 +6,28 @@
     Chain,
     Inline,
     Fixpoint,
-    WrapConst,
     Call2Invoke,
     ConstantFold,
     CFGCompactify,
     InlineGetItem,
     InlineGetField,
     DeadCodeElimination,
 )
-from kirin.analysis import const
 from kirin.ir.method import Method
 from kirin.rewrite.abc import RewriteResult
+from kirin.passes.hint_const import HintConst
 
 
 @dataclass
 class Fold(Pass):
-    constprop: const.Propagate = field(init=False)
+    hint_const: HintConst = field(init=False)
 
     def __post_init__(self):
-        self.constprop = const.Propagate(self.dialects)
+        self.hint_const = HintConst(self.dialects)
+        self.hint_const.no_raise = self.no_raise
 
     def unsafe_run(self, mt: Method) -> RewriteResult:
-        result = RewriteResult()
-        frame, _ = self.constprop.run_analysis(mt, no_raise=self.no_raise)
-        result = Walk(WrapConst(frame)).rewrite(mt.code).join(result)
+        result = self.hint_const.unsafe_run(mt)
         rule = Chain(
             ConstantFold(),
             Call2Invoke(),

src/kirin/passes/default.py

@@ -0,0 +1,49 @@
+from dataclasses import field, dataclass
+
+from kirin.dialects import ilist
+from kirin.ir.method import Method
+from kirin.passes.fold import Fold
+from kirin.rewrite.abc import RewriteResult
+from kirin.passes.aggressive import Fold as AggressiveFold
+
+from .abc import Pass
+from .typeinfer import TypeInfer
+from .hint_const import HintConst
+
+
+@dataclass
+class Default(Pass):
+    verify: bool = field(default=False, kw_only=True)
+    fold: bool = field(default=True, kw_only=True)
+    aggressive: bool = field(default=False, kw_only=True)
+    typeinfer: bool = field(default=True, kw_only=True)
+
+    hint_const_pass: HintConst = field(init=False)
+    typeinfer_pass: TypeInfer = field(init=False)
+    ilist_desugar: ilist.IListDesugar = field(init=False)
+    fold_pass: Pass = field(init=False)
+
+    def __post_init__(self):
+        # TODO: cleanup no_raise
+        self.ilist_desugar = ilist.IListDesugar(self.dialects, no_raise=self.no_raise)
+        self.typeinfer_pass = TypeInfer(self.dialects, no_raise=self.no_raise)
+        self.hint_const_pass = HintConst(self.dialects, no_raise=self.no_raise)
+        if self.aggressive:
+            self.fold_pass = AggressiveFold(self.dialects, no_raise=self.no_raise)
+        else:
+            self.fold_pass = Fold(self.dialects, no_raise=self.no_raise)
+
+    def unsafe_run(self, mt: Method) -> RewriteResult:
+        if self.verify:
+            mt.verify()
+
+        result = self.ilist_desugar.fixpoint(mt)
+        if self.typeinfer:
+            self.typeinfer_pass(mt).join(result)
+
+        if self.fold:
+            if self.aggressive:
+                self.fold_pass.fixpoint(mt).join(result)
+            else:
+                self.fold_pass(mt).join(result)
+        return result

src/kirin/passes/fold.py

@@ -1,29 +1,32 @@
-from dataclasses import dataclass
+from dataclasses import field, dataclass
 
 from kirin.ir import Method, SSACFGRegion
 from kirin.rewrite import (
     Walk,
     Chain,
     Fixpoint,
-    WrapConst,
     Call2Invoke,
     ConstantFold,
     CFGCompactify,
     InlineGetItem,
     DeadCodeElimination,
 )
-from kirin.analysis import const
 from kirin.passes.abc import Pass
 from kirin.rewrite.abc import RewriteResult
 
+from .hint_const import HintConst
+
 
 @dataclass
 class Fold(Pass):
+    hint_const: HintConst = field(init=False)
+
+    def __post_init__(self):
+        self.hint_const = HintConst(self.dialects)
+        self.hint_const.no_raise = self.no_raise
 
     def unsafe_run(self, mt: Method) -> RewriteResult:
-        constprop = const.Propagate(self.dialects)
-        frame, _ = constprop.run_analysis(mt, no_raise=self.no_raise)
-        result = Walk(WrapConst(frame)).rewrite(mt.code)
+        result = self.hint_const.unsafe_run(mt)
         result = (
             Fixpoint(
                 Walk(

src/kirin/passes/hint_const.py

@@ -0,0 +1,16 @@
+from dataclasses import dataclass
+
+from kirin.ir import Method
+from kirin.rewrite import Walk, WrapConst
+from kirin.analysis import const
+from kirin.passes.abc import Pass
+from kirin.rewrite.abc import RewriteResult
+
+
+@dataclass
+class HintConst(Pass):
+
+    def unsafe_run(self, mt: Method) -> RewriteResult:
+        constprop = const.Propagate(self.dialects)
+        frame, _ = constprop.run_analysis(mt, no_raise=self.no_raise)
+        return Walk(WrapConst(frame)).rewrite(mt.code)

src/kirin/passes/typeinfer.py

@@ -1,4 +1,4 @@
-from dataclasses import dataclass
+from dataclasses import field, dataclass
 
 from kirin.ir import Method, HasSignature
 from kirin.rewrite import Walk, Chain
@@ -8,23 +8,35 @@
 from kirin.analysis.typeinfer import TypeInference
 from kirin.rewrite.apply_type import ApplyType
 from kirin.rewrite.type_assert import InlineTypeAssert
+from kirin.dialects.ilist.rewrite import HintLen
+
+from .hint_const import HintConst
 
 
 @dataclass
 class TypeInfer(Pass):
+    hint_const: HintConst = field(init=False)
 
     def __post_init__(self):
         self.infer = TypeInference(self.dialects)
+        self.hint_const = HintConst(self.dialects)
+        self.hint_const.no_raise = self.no_raise
 
     def unsafe_run(self, mt: Method) -> RewriteResult:
+        result = self.hint_const.unsafe_run(mt)
         frame, return_type = self.infer.run_analysis(
             mt, mt.arg_types, no_raise=self.no_raise
         )
         if trait := mt.code.get_trait(HasSignature):
             trait.set_signature(mt.code, Signature(mt.arg_types, return_type))
 
-        result = Chain(
-            Walk(ApplyType(frame.entries)), Walk(InlineTypeAssert())
-        ).rewrite(mt.code)
+        result = (
+            Chain(
+                Walk(ApplyType(frame.entries)),
+                Walk(Chain(InlineTypeAssert(), HintLen())),
+            )
+            .rewrite(mt.code)
+            .join(result)
+        )
         mt.inferred = True
         return result

src/kirin/prelude.py

@@ -6,7 +6,7 @@
 from typing_extensions import Doc, Annotated
 
 from kirin.ir import Method, dialect_group
-from kirin.passes import aggressive
+from kirin.passes import Default
 from kirin.dialects import cf, scf, func, math, ilist, lowering
 from kirin.dialects.py import (
     cmp,
@@ -26,8 +26,6 @@
     iterable,
     assertion,
 )
-from kirin.passes.fold import Fold
-from kirin.passes.typeinfer import TypeInfer
 
 
 @dialect_group(
@@ -142,10 +140,6 @@ def main(x: int) -> int:
     main.print() # main is a Method!
     ```
     """
-    fold_pass = Fold(self)
-    ilist_desugar = ilist.IListDesugar(self)
-    aggressive_fold_pass = aggressive.Fold(self)
-    typeinfer_pass = TypeInfer(self)
 
     def run_pass(
         mt: Annotated[Method, Doc("The method to run pass on.")],
@@ -163,20 +157,17 @@ def run_pass(
         aggressive: Annotated[
             bool, Doc("run aggressive folding passes if `fold=True`")
         ] = False,
+        no_raise: Annotated[bool, Doc("do not raise exception during analysis")] = True,
     ) -> None:
-        if verify:
-            mt.verify()
-
-        ilist_desugar.fixpoint(mt)
-
-        if fold:
-            if aggressive:
-                aggressive_fold_pass.fixpoint(mt)
-            else:
-                fold_pass(mt)
-
-        if typeinfer:
-            typeinfer_pass(mt)
+        default_pass = Default(
+            self,
+            verify=verify,
+            fold=fold,
+            aggressive=aggressive,
+            typeinfer=typeinfer,
+            no_raise=no_raise,
+        )
+        default_pass.fixpoint(mt)
 
     return run_pass
 
@@ -221,16 +212,34 @@ def run_pass(method: Method) -> None:
     )
 )
 def structural(self):
-    """Structural kernel without optimization passes."""
-    typeinfer_pass = TypeInfer(self)
+    """Structural kernel with optimization passes."""
 
     def run_pass(
-        method: Method, *, verify: bool = True, typeinfer: bool = True
+        mt: Annotated[Method, Doc("The method to run pass on.")],
+        *,
+        verify: Annotated[
+            bool, Doc("run `verify` before running passes, default is `True`")
+        ] = True,
+        typeinfer: Annotated[
+            bool,
+            Doc(
+                "run type inference and apply the inferred type to IR, default `False`"
+            ),
+        ] = False,
+        fold: Annotated[bool, Doc("run folding passes")] = True,
+        aggressive: Annotated[
+            bool, Doc("run aggressive folding passes if `fold=True`")
+        ] = False,
+        no_raise: Annotated[bool, Doc("do not raise exception during analysis")] = True,
     ) -> None:
-        if verify:
-            method.verify()
-
-        if typeinfer:
-            typeinfer_pass(method)
+        default_pass = Default(
+            self,
+            verify=verify,
+            fold=fold,
+            aggressive=aggressive,
+            typeinfer=typeinfer,
+            no_raise=no_raise,
+        )
+        default_pass.fixpoint(mt)
 
     return run_pass