From 68bf8111094c5fd0f6d6b9f36081fb8708cf8bcb Mon Sep 17 00:00:00 2001 From: John Long Date: Fri, 3 Apr 2026 14:43:39 -0400 Subject: [PATCH 1/6] carry early termination optimization to constprop --- src/kirin/dialects/scf/constprop.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/src/kirin/dialects/scf/constprop.py b/src/kirin/dialects/scf/constprop.py index 25181769c..9dda110fa 100644 --- a/src/kirin/dialects/scf/constprop.py +++ b/src/kirin/dialects/scf/constprop.py @@ -108,6 +108,7 @@ def _prop_const_iterable_forloop( loop_vars = frame.get_values(stmt.initializers) + prev_loop_vars = None for value in iterable.data: with interp_.new_frame(stmt, has_parent_access=True) as body_frame: loop_vars = interp_.frame_call_region( @@ -121,6 +122,10 @@ def _prop_const_iterable_forloop( elif isinstance(loop_vars, interp.ReturnValue): return loop_vars + if prev_loop_vars is not None and loop_vars == prev_loop_vars: + break + prev_loop_vars = loop_vars + if not frame_is_not_pure: frame.should_be_pure.add(stmt) return loop_vars From d9474985ab5826885d59048df9015f38d767db4c Mon Sep 17 00:00:00 2001 From: John Long Date: Fri, 3 Apr 2026 15:17:34 -0400 Subject: [PATCH 2/6] do not allow early termination if purity cannot be guaranteed --- src/kirin/dialects/scf/constprop.py | 17 ++++++++++- test/dialects/scf/test_constprop.py | 46 +++++++++++++++++++++++++++++ 2 files changed, 62 insertions(+), 1 deletion(-) diff --git a/src/kirin/dialects/scf/constprop.py b/src/kirin/dialects/scf/constprop.py index 9dda110fa..e4a1f2407 100644 --- a/src/kirin/dialects/scf/constprop.py +++ b/src/kirin/dialects/scf/constprop.py @@ -108,6 +108,17 @@ def _prop_const_iterable_forloop( loop_vars = frame.get_values(stmt.initializers) + # Early termination is only sound when the loop body does not + # reference the iteration variable. When the body ignores it, + # each iteration is a pure function of loop_vars alone, so + # convergence of loop_vars implies all subsequent iterations + # are identical (including purity). If the body *does* use the + # iteration variable, different iterations may follow different + # code paths (e.g. conditional side effects on a particular i), + # so we must execute every iteration. + iter_var = stmt.body.blocks[0].args[0] + can_early_terminate = not iter_var.uses + prev_loop_vars = None for value in iterable.data: with interp_.new_frame(stmt, has_parent_access=True) as body_frame: @@ -122,7 +133,11 @@ def _prop_const_iterable_forloop( elif isinstance(loop_vars, interp.ReturnValue): return loop_vars - if prev_loop_vars is not None and loop_vars == prev_loop_vars: + if ( + can_early_terminate + and prev_loop_vars is not None + and loop_vars == prev_loop_vars + ): break prev_loop_vars = loop_vars diff --git a/test/dialects/scf/test_constprop.py b/test/dialects/scf/test_constprop.py index 30d3535b3..738e5423d 100644 --- a/test/dialects/scf/test_constprop.py +++ b/test/dialects/scf/test_constprop.py @@ -1,11 +1,22 @@ from pytest import mark +from kirin import ir, lowering +from kirin.decl import statement from kirin.prelude import structural_no_opt from kirin.analysis import const from kirin.dialects import scf, func prop = const.Propagate(structural_no_opt) +# A statement with no Pure/MaybePure trait — acts as a side effect. +_impure_dialect = ir.Dialect("test_impure") + + +@statement(dialect=_impure_dialect) +class ImpureOp(ir.Statement): + name = "impure_op" + traits = frozenset({lowering.FromPythonCall()}) + def test_simple_loop(): @structural_no_opt @@ -103,3 +114,38 @@ def simple_ifelse(x: int): assert isinstance(terminator, func.Return) assert isinstance(value := frame.entries[terminator.value], const.Value) assert value.data == 0 + + +def test_no_early_termination_when_body_uses_iter_var(): + """Early termination must not fire when the body references the iteration + variable, because later iterations may follow different code paths that + affect purity. Here the impure ``ImpureOp`` is guarded by ``i == 2``, + so the loop body is impure only on iteration 2. If early termination + incorrectly broke after iteration 1 (where loop_vars converge), the + for-loop would be marked as pure when it is not.""" + + _group = structural_no_opt.add(_impure_dialect) + + @_group + def impure_on_later_iter(x: int) -> int: + for i in range(5): + if i == 2: + ImpureOp() + x = x + 1 + return x + + constprop = const.Propagate(_group) + frame, ret = constprop.run(impure_on_later_iter) + + # The for-loop statement is in the first block of the callable region. + for_stmt = None + for block in impure_on_later_iter.callable_region.blocks: + for stmt in block.stmts: + if isinstance(stmt, scf.For): + for_stmt = stmt + break + + assert for_stmt is not None, "Could not find scf.For in the IR" + # The for-loop must NOT be in should_be_pure — it contains a + # conditionally-impure operation on a later iteration. + assert for_stmt not in frame.should_be_pure From c5dcb51c8a464b473885e18386faf18ef62ddd17 Mon Sep 17 00:00:00 2001 From: John Long Date: Fri, 3 Apr 2026 15:24:21 -0400 Subject: [PATCH 3/6] simplify early termination comment Co-Authored-By: Claude Opus 4.6 (1M context) --- src/kirin/dialects/scf/constprop.py | 10 ++-------- 1 file changed, 2 insertions(+), 8 deletions(-) diff --git a/src/kirin/dialects/scf/constprop.py b/src/kirin/dialects/scf/constprop.py index e4a1f2407..11303718c 100644 --- a/src/kirin/dialects/scf/constprop.py +++ b/src/kirin/dialects/scf/constprop.py @@ -108,14 +108,8 @@ def _prop_const_iterable_forloop( loop_vars = frame.get_values(stmt.initializers) - # Early termination is only sound when the loop body does not - # reference the iteration variable. When the body ignores it, - # each iteration is a pure function of loop_vars alone, so - # convergence of loop_vars implies all subsequent iterations - # are identical (including purity). If the body *does* use the - # iteration variable, different iterations may follow different - # code paths (e.g. conditional side effects on a particular i), - # so we must execute every iteration. + # Only safe to break early when the body doesn't use the iteration + # variable — otherwise later iterations may take different code paths. iter_var = stmt.body.blocks[0].args[0] can_early_terminate = not iter_var.uses From 074afa19bd0fd3ca2c1eda04b061e6aa825cb591 Mon Sep 17 00:00:00 2001 From: John Long Date: Fri, 3 Apr 2026 15:31:16 -0400 Subject: [PATCH 4/6] simplify scf.For lookup in early termination test Co-Authored-By: Claude Opus 4.6 (1M context) --- test/dialects/scf/test_constprop.py | 10 +--------- 1 file changed, 1 insertion(+), 9 deletions(-) diff --git a/test/dialects/scf/test_constprop.py b/test/dialects/scf/test_constprop.py index 738e5423d..70d6b3245 100644 --- a/test/dialects/scf/test_constprop.py +++ b/test/dialects/scf/test_constprop.py @@ -137,15 +137,7 @@ def impure_on_later_iter(x: int) -> int: constprop = const.Propagate(_group) frame, ret = constprop.run(impure_on_later_iter) - # The for-loop statement is in the first block of the callable region. - for_stmt = None - for block in impure_on_later_iter.callable_region.blocks: - for stmt in block.stmts: - if isinstance(stmt, scf.For): - for_stmt = stmt - break - - assert for_stmt is not None, "Could not find scf.For in the IR" + [for_stmt] = [s for s in impure_on_later_iter.code.walk() if isinstance(s, scf.For)] # The for-loop must NOT be in should_be_pure — it contains a # conditionally-impure operation on a later iteration. assert for_stmt not in frame.should_be_pure From aba74210d3ebc5a7577c1076fd329dcbd01f1703 Mon Sep 17 00:00:00 2001 From: John Long Date: Fri, 3 Apr 2026 15:36:29 -0400 Subject: [PATCH 5/6] increase iteration count in early termination safety test Co-Authored-By: Claude Opus 4.6 (1M context) --- test/dialects/scf/test_constprop.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/test/dialects/scf/test_constprop.py b/test/dialects/scf/test_constprop.py index 70d6b3245..cd46ff2ba 100644 --- a/test/dialects/scf/test_constprop.py +++ b/test/dialects/scf/test_constprop.py @@ -119,8 +119,8 @@ def simple_ifelse(x: int): def test_no_early_termination_when_body_uses_iter_var(): """Early termination must not fire when the body references the iteration variable, because later iterations may follow different code paths that - affect purity. Here the impure ``ImpureOp`` is guarded by ``i == 2``, - so the loop body is impure only on iteration 2. If early termination + affect purity. Here the impure ``ImpureOp`` is guarded by ``i == 50``, + so the loop body is impure only on iteration 50. If early termination incorrectly broke after iteration 1 (where loop_vars converge), the for-loop would be marked as pure when it is not.""" @@ -128,8 +128,8 @@ def test_no_early_termination_when_body_uses_iter_var(): @_group def impure_on_later_iter(x: int) -> int: - for i in range(5): - if i == 2: + for i in range(100): + if i == 50: ImpureOp() x = x + 1 return x From 582bd7ff849a638f46cbccc27bacd790120e8d17 Mon Sep 17 00:00:00 2001 From: John Long Date: Fri, 3 Apr 2026 15:49:50 -0400 Subject: [PATCH 6/6] add happy path test for constprop early termination Co-Authored-By: Claude Opus 4.6 (1M context) --- test/dialects/scf/test_constprop.py | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/test/dialects/scf/test_constprop.py b/test/dialects/scf/test_constprop.py index cd46ff2ba..875ee5556 100644 --- a/test/dialects/scf/test_constprop.py +++ b/test/dialects/scf/test_constprop.py @@ -116,6 +116,25 @@ def simple_ifelse(x: int): assert value.data == 0 +def test_early_termination_when_body_ignores_iter_var(): + """When the body doesn't reference the iteration variable and loop_vars + converge (x is Unknown, so Unknown + 1 = Unknown), early termination + should fire and produce the same result as running all iterations.""" + + @structural_no_opt + def converging_loop(x: int) -> int: + for _i in range(100): + x = x + 1 + return x + + constprop = const.Propagate(structural_no_opt) + frame, ret = constprop.run(converging_loop) + + assert isinstance(ret, const.Unknown) + [for_stmt] = [s for s in converging_loop.code.walk() if isinstance(s, scf.For)] + assert for_stmt in frame.should_be_pure + + def test_no_early_termination_when_body_uses_iter_var(): """Early termination must not fire when the body references the iteration variable, because later iterations may follow different code paths that