diff --git a/src/kirin/analysis/typeinfer/widen.py b/src/kirin/analysis/typeinfer/widen.py new file mode 100644 index 0000000000..f67c9f28e7 --- /dev/null +++ b/src/kirin/analysis/typeinfer/widen.py @@ -0,0 +1,26 @@ +from kirin import types + + +def widen( + previous: types.TypeAttribute, current: types.TypeAttribute +) -> types.TypeAttribute: + """Include both types, forgetting changing literals within matching generics.""" + if current.is_subseteq(previous): + return previous + if previous.is_subseteq(types.Bottom): + return current + if ( + isinstance(previous, types.Generic) + and isinstance(current, types.Generic) + and previous.body == current.body + and len(previous.vars) == len(current.vars) + and previous.vararg is None + and current.vararg is None + ): + return types.Generic( + previous.body, + *(widen(old, new) for old, new in zip(previous.vars, current.vars)), + ) + if isinstance(previous, types.Literal) and isinstance(current, types.Literal): + return types.Any + return previous.join(current) diff --git a/src/kirin/dialects/ilist/typeinfer.py b/src/kirin/dialects/ilist/typeinfer.py index 47c1f099db..aa69f019b2 100644 --- a/src/kirin/dialects/ilist/typeinfer.py +++ b/src/kirin/dialects/ilist/typeinfer.py @@ -43,7 +43,7 @@ def eltype_list( def new(self, interp: TypeInference, frame: Frame[types.TypeAttribute], stmt: New): values = frame.get_values(stmt.values) if not values: - return (IListType[types.Any, types.Literal(0)],) + return (IListType[stmt.elem_type, types.Literal(0)],) elem_type = values[0] for v in values: diff --git a/src/kirin/dialects/scf/typeinfer.py b/src/kirin/dialects/scf/typeinfer.py index 7cb33f35f0..ddaa10ca46 100644 --- a/src/kirin/dialects/scf/typeinfer.py +++ b/src/kirin/dialects/scf/typeinfer.py @@ -2,11 +2,16 @@ from kirin.analysis import ForwardFrame, TypeInference from kirin.dialects import func from kirin.dialects.eltype import ElType +from kirin.analysis.typeinfer.widen import widen from . import absint from .stmts import For, IfElse from ._dialect import dialect +# Precision/performance heuristic: allow structural widening before falling back +# to top. Correctness requires a verified invariant regardless of this budget. +_LOOP_WIDENING_BUDGET = 16 + @dialect.register(key="typeinfer") class TypeInfer(absint.Methods): @@ -40,16 +45,48 @@ def for_loop( if not isinstance(eltype, tuple): # error return item = eltype[0] - frame.set_values(block_args, (item,) + loop_vars) if isinstance(body_block.last_stmt, func.Return): + frame.set_values(block_args, (item,) + loop_vars) frame.worklist.append(interp.Successor(body_block, item, *loop_vars)) return # if terminate is Return, there is no result - loop_vars_ = interp_.frame_call_region(frame, stmt, stmt.body, item, *loop_vars) - if isinstance(loop_vars_, interp.ReturnValue): - return loop_vars_ - elif isinstance(loop_vars_, tuple): - return interp_.join_results(loop_vars, loop_vars_) - else: # None, loop has no result - return + candidate = loop_vars + # Each fallback update promotes at least one changing component to top. + # Allow one update per component and a final invariant check. + max_iterations = _LOOP_WIDENING_BUDGET + len(candidate) + 1 + for iteration in range(max_iterations): + # A fresh frame prevents earlier visits and intermediate types from + # contaminating inference under the current loop invariant. + with interp_.new_frame(stmt, has_parent_access=True) as body_frame: + yielded = interp_.frame_call_region( + body_frame, stmt, stmt.body, item, *candidate + ) + + if not isinstance(yielded, tuple): + frame.set_values(body_frame.entries.keys(), body_frame.entries.values()) + return yielded + + required = tuple( + initial.join(value) for initial, value in zip(loop_vars, yielded) + ) + if all(new.is_subseteq(old) for old, new in zip(candidate, required)): + frame.set_values(body_frame.entries.keys(), body_frame.entries.values()) + frame.set_values(block_args, (item,) + candidate) + return candidate + + # Structural growth can also produce unbounded unions or nesting. + # Fall back to top for changing components, then verify the invariant + # with another body evaluation rather than returning a partial result. + candidate = tuple( + ( + widen(old, new) + if iteration < _LOOP_WIDENING_BUDGET + else old if new.is_subseteq(old) else types.Any + ) + for old, new in zip(candidate, yielded) + ) + + raise interp.InterpreterError( + f"scf.For type inference did not converge after {max_iterations} iterations" + ) diff --git a/test/analysis/dataflow/typeinfer/test_widen.py b/test/analysis/dataflow/typeinfer/test_widen.py new file mode 100644 index 0000000000..e2259872fa --- /dev/null +++ b/test/analysis/dataflow/typeinfer/test_widen.py @@ -0,0 +1,36 @@ +from pytest import mark + +from kirin import types +from kirin.dialects.ilist import IListType +from kirin.analysis.typeinfer.widen import widen + + +@mark.parametrize( + "previous,current,expected", + [ + (types.Bottom, types.Int, types.Int), + (types.Any, types.Int, types.Any), + (types.Int, types.Float, types.Int | types.Float), + ( + IListType[types.Int, types.Literal(1)], + IListType[types.Float, types.Literal(1)], + IListType[types.Int | types.Float, types.Literal(1)], + ), + ( + IListType[types.Int, types.Literal(0)], + IListType[types.Int, types.Literal(1)], + IListType[types.Int, types.Any], + ), + ( + types.Tuple[IListType[types.Int, types.Literal(0)], types.Bool], + types.Tuple[IListType[types.Int, types.Literal(1)], types.Bool], + types.Tuple[IListType[types.Int, types.Any], types.Bool], + ), + ], +) +def test_widen_preserves_upper_bounds(previous, current, expected): + result = widen(previous, current) + assert result == expected + assert previous.is_subseteq(result) + assert current.is_subseteq(result) + assert widen(result, current) == result diff --git a/test/dialects/scf/test_loop_inference.py b/test/dialects/scf/test_loop_inference.py new file mode 100644 index 0000000000..02029f7a87 --- /dev/null +++ b/test/dialects/scf/test_loop_inference.py @@ -0,0 +1,249 @@ +from typing import Any + +from pytest import mark + +from kirin import types +from kirin.passes import TypeInfer +from kirin.prelude import structural_no_opt +from kirin.rewrite import Walk +from kirin.analysis import TypeInference +from kirin.dialects import scf, ilist + + +@structural_no_opt +def identity(x: int) -> int: + return x + + +@mark.parametrize("n", [0, 1, 3]) +def test_unroll_map_in_growing_loop(n): + @structural_no_opt + def grow(n: int) -> ilist.IList[int, Any]: + xs = ilist.IList([]) + for _ in range(n): + xs = xs + [4] + xs = ilist.map(identity, xs) + return xs + + expected = list(grow(n)) + TypeInfer(grow.dialects, no_raise=False)(grow) + Walk(ilist.rewrite.Unroll()).rewrite(grow.code) + assert list(grow(n)) == expected + + +def test_growing_list_invariant(): + @structural_no_opt + def grow(n: int): + xs = [1] + for i in range(n): + xs = xs + [i] + return xs + + frame, result = TypeInference(grow.dialects).run(grow) + expected = ilist.IListType[types.Int, types.Any] + assert result == expected + loop = next(stmt for stmt in grow.code.walk() if isinstance(stmt, scf.For)) + assert all(frame.get(arg) == expected for arg in loop.body.blocks[0].args[1:]) + + +def test_stable_length_and_multiple_carried_values(): + @structural_no_opt + def loop(n: int): + fixed = [1] + growing = [2] + for i in range(n): + fixed = [i] + growing = growing + fixed + return fixed, growing + + _, result = TypeInference(loop.dialects).run(loop) + assert ( + result + == types.Tuple[ + ilist.IListType[types.Int, types.Literal(1)], + ilist.IListType[types.Int, types.Any], + ] + ) + + +def test_nested_loops_and_repeated_inference(): + @structural_no_opt + def grow(n: int): + xs = [1] + for i in range(n): + for j in range(n): + xs = xs + [i + j] + return xs + + expected = list(grow(3)) + inference = TypeInfer(grow.dialects, no_raise=False) + inference(grow) + assert grow.return_type == ilist.IListType[types.Int, types.Any] + first = [value.type for stmt in grow.code.walk() for value in stmt.results] + inference(grow) + assert [value.type for stmt in grow.code.walk() for value in stmt.results] == first + assert list(grow(3)) == expected + + +def test_changing_element_type(): + @structural_no_opt + def grow(n: int): + xs = [1] + for _ in range(n): + xs = xs + [1.5] + return xs + + _, result = TypeInference(grow.dialects).run(grow) + assert result == ilist.IListType[types.Int | types.Float, types.Any] + + +def test_nested_tuple_growth_terminates(): + @structural_no_opt + def nest(n: int): + value = () + for _ in range(n): + value = (value,) + return value + + _, result = TypeInference(nest.dialects).run(nest) + assert result == types.Any + + +def test_typed_empty_list(): + from kirin.interp import Frame + from kirin.dialects.ilist.typeinfer import TypeInfer as IListTypeInfer + + stmt = ilist.New((), elem_type=types.Int) + result = IListTypeInfer().new(TypeInference(structural_no_opt), Frame(stmt), stmt) + assert result == (ilist.IListType[types.Int, types.Literal(0)],) + + +def test_inconsistent_lattice_fails_with_bounded_iterations(): + from contextlib import nullcontext + from unittest.mock import Mock + + from pytest import raises + + from kirin import interp + from kirin.analysis import ForwardFrame + from kirin.dialects.scf.typeinfer import TypeInfer as SCFTypeInfer + + @structural_no_opt + def grow(n: int): + xs = [1] + for i in range(n): + xs = xs + [i] + return xs + + loop = next(stmt for stmt in grow.code.walk() if isinstance(stmt, scf.For)) + # Model a broken lattice that rejects containment even after widening to top. + value = Mock(spec=types.TypeAttribute) + value.is_subseteq.return_value = False + value.join.return_value = value + values = (value,) * len(loop.initializers) + frame = Mock(spec=ForwardFrame) + frame.get_values.return_value = values + interpreter = Mock(spec=TypeInference) + interpreter.frame_eval.return_value = (types.Int,) + interpreter.frame_call_region.return_value = values + interpreter.new_frame.side_effect = lambda *args, **kwargs: nullcontext( + ForwardFrame(loop) + ) + + with raises( + interp.InterpreterError, match="scf.For type inference did not converge" + ): + SCFTypeInfer().for_loop(interpreter, frame, loop) + + assert interpreter.frame_call_region.call_count <= 32 + frame.set_values.assert_not_called() + + +@mark.parametrize("n", [0, 1, 2, 5]) +def test_branch_dependent_growth_preserves_execution(n): + @structural_no_opt + def grow(n: int): + xs = [1] + for i in range(n): + if i % 2 == 0: + xs = xs + [i] + else: + xs = xs + [i, i] + xs = ilist.map(identity, xs) + return xs + + expected = list(grow(n)) + TypeInfer(grow.dialects, no_raise=False)(grow) + maps = [stmt for stmt in grow.code.walk() if isinstance(stmt, ilist.Map)] + assert len(maps) == 1 + assert ilist.IListType[types.Int, types.Any].is_subseteq(maps[0].collection.type) + Walk(ilist.rewrite.Unroll()).rewrite(grow.code) + assert list(grow(n)) == expected + + +@mark.parametrize("n", [0, 1, 2, 5]) +def test_dependent_carried_lists_preserve_execution(n): + @structural_no_opt + def grow(n: int): + xs = [1] + ys = [2] + for i in range(n): + xs = ilist.map(identity, ys) + ys = xs + [i] + return xs, ys + + before_xs, before_ys = grow(n) + TypeInfer(grow.dialects, no_raise=False)(grow) + mapped = next(stmt for stmt in grow.code.walk() if isinstance(stmt, ilist.Map)) + assert mapped.collection.type == ilist.IListType[types.Int, types.Any] + Walk(ilist.rewrite.Unroll()).rewrite(grow.code) + after_xs, after_ys = grow(n) + assert list(after_xs) == list(before_xs) + assert list(after_ys) == list(before_ys) + + +@mark.parametrize("budget", [0, 1, 16]) +def test_budget_fallback_preserves_stable_components(monkeypatch, budget): + monkeypatch.setattr("kirin.dialects.scf.typeinfer._LOOP_WIDENING_BUDGET", budget) + + @structural_no_opt + def nest(n: int): + value = () + fixed = [1] + for i in range(n): + value = (value,) + fixed = [i] + return value, fixed + + before_value, before_fixed = nest(3) + TypeInfer(nest.dialects, no_raise=False)(nest) + assert ( + nest.return_type + == types.Tuple[types.Any, ilist.IListType[types.Int, types.Literal(1)]] + ) + Walk(ilist.rewrite.Unroll()).rewrite(nest.code) + after_value, after_fixed = nest(3) + assert after_value == before_value + assert list(after_fixed) == list(before_fixed) + + +def test_fallback_propagates_through_dependent_values(monkeypatch): + monkeypatch.setattr("kirin.dialects.scf.typeinfer._LOOP_WIDENING_BUDGET", 0) + + @structural_no_opt + def nest(n: int): + first = () + second = () + third = () + for _ in range(n): + first = second + second = third + third = (third,) + return first, second, third + + expected = [nest(n) for n in (0, 1, 4)] + TypeInfer(nest.dialects, no_raise=False)(nest) + assert nest.return_type == types.Tuple[types.Any, types.Any, types.Any] + loop = next(stmt for stmt in nest.code.walk() if isinstance(stmt, scf.For)) + assert all(arg.type == types.Any for arg in loop.body.blocks[0].args[1:]) + assert [nest(n) for n in (0, 1, 4)] == expected