diff --git a/ptodsl/ptodsl/_ast_rewrite.py b/ptodsl/ptodsl/_ast_rewrite.py index 36cd427bd6..7714f251b3 100644 --- a/ptodsl/ptodsl/_ast_rewrite.py +++ b/ptodsl/ptodsl/_ast_rewrite.py @@ -601,6 +601,19 @@ def _slot_live_before_block(stmts, live_after, static_env, static_iters=None) -> def _slot_live_before_stmt(stmt, live_after, static_env, static_iters) -> set[_SubscriptSlot]: + if isinstance(stmt, (ast.With, ast.AsyncWith)): + live = _slot_live_before_block(stmt.body, live_after, static_env, static_iters) + # Python evaluates with-items from left to right and binds each + # optional_vars immediately. Reverse the sequence for liveness so a + # context expression can use a binding produced by an earlier item + # without incorrectly turning that use into a live-in. + for item in reversed(stmt.items): + if item.optional_vars is not None: + live = _kill_slots_for_with_target( + live, item.optional_vars, static_env, static_iters + ) + live |= _slot_info(item.context_expr, static_env, static_iters).loads + return live if isinstance(stmt, ast.If): test_info = _slot_info(stmt.test, static_env, static_iters) return ( @@ -648,6 +661,37 @@ def _kill_slots_for_assigned_bases(slots, stmt) -> set[_SubscriptSlot]: } +def _kill_slots_for_with_target( + slots, target, static_env, static_iters +) -> set[_SubscriptSlot]: + target_info = _slot_info(target, static_env, static_iters) + bound_bases = _simple_name_targets(target) + dynamic_subscript_bases = set() + for subscript in _target_subscripts(target): + if not _resolve_subscript_slots( + subscript, static_env, static_iters, require_static=True + ) and isinstance(subscript.value, ast.Name): + dynamic_subscript_bases.add(subscript.value.id) + killed_bases = bound_bases | dynamic_subscript_bases + return { + slot + for slot in slots + if slot.base not in killed_bases and slot not in target_info.stores + } + + +def _target_subscripts(target): + if isinstance(target, ast.Subscript): + yield target + return + if isinstance(target, (ast.Tuple, ast.List)): + for element in target.elts: + yield from _target_subscripts(element) + return + if isinstance(target, ast.Starred): + yield from _target_subscripts(target.value) + + def _assigned_name_targets(stmt) -> set[str]: if isinstance(stmt, ast.Assign): names = set() diff --git a/ptodsl/tests/test_jit_compile.py b/ptodsl/tests/test_jit_compile.py index ae7b682536..95ba2ac878 100644 --- a/ptodsl/tests/test_jit_compile.py +++ b/ptodsl/tests/test_jit_compile.py @@ -516,6 +516,39 @@ def _assert_ast_rewrite_nested_partial_assign_ssa_identity(): 'innermost else of the slot probe must not yield the sibling branch q2 value', ) + # Focused probe 4: with-items bind their optional_vars in source order. + # The second context expression below reads the value produced by the first + # item, so it must not become a live-in slot. A subscript target also kills + # only the slot it overwrites, rather than the whole list binding. + import ast as _slot_ast + from ptodsl._ast_rewrite import _SubscriptSlot, _read_before_assignment_slots + + def with_slot_liveness(source, live_after=()): + with_stmt = _slot_ast.parse(source).body[0] + return _read_before_assignment_slots( + [with_stmt], {}, live_after=set(live_after) + ) + + values0 = _SubscriptSlot('values', 0) + values1 = _SubscriptSlot('values', 1) + expect( + not with_slot_liveness( + 'with cm() as values, cm(values[0]):\n pass' + ), + 'a later with-item may use the binding produced by an earlier item', + ) + expect( + with_slot_liveness('with cm(values[0]) as values:\n pass') == {values0}, + 'a with context expression must still keep an outer slot live-in', + ) + expect( + with_slot_liveness( + 'with cm() as values[0]:\n pass', + live_after={values0, values1}, + ) == {values1}, + 'a subscript with-as target must kill only the overwritten slot', + ) + # Augmenting the induction variable must not invent a carry for it. iv_augassign_text = ast_for_iv_augassign_probe.compile().mlir_text() expect_parse_roundtrip_and_verify(iv_augassign_text, 'AST-rewritten induction-variable augment-assignment') @@ -2117,6 +2150,22 @@ def ast_runtime_for_static_slot_carry_probe(cols: pto.i32): _ = total +@pto.jit(target="a5", mode="explicit") +def ast_runtime_for_static_slot_nested_vecscope_probe(cols: pto.i32): + zero = pto.const(0, dtype=pto.index) + + for _ in range(cols): + with pto.vecscope(): + values = [zero, zero, zero, zero] + for c in range(64): + values[0] = c + values[1] = c + values[2] = c + values[3] = c + total = values[0] + values[1] + values[2] + values[3] + _ = total + + _STATIC_SLOT_GLOBAL_INDEX = 2 @@ -7107,6 +7156,17 @@ def _enter_inline_simt_with_resource_attr(): and "scf.yield" in ast_runtime_for_static_slot_carry_text, "static subscript slot carry should lower through scf.for iter_args", ) + ast_runtime_for_static_slot_nested_vecscope_text = ( + ast_runtime_for_static_slot_nested_vecscope_probe.compile().mlir_text() + ) + expect_parse_roundtrip_and_verify( + ast_runtime_for_static_slot_nested_vecscope_text, + "AST-rewritten nested vecscope runtime for static subscript slots", + ) + expect( + ast_runtime_for_static_slot_nested_vecscope_text.count("scf.for") == 2, + "nested vecscope static subscript slots should preserve both authored runtime loops", + ) ast_runtime_for_static_slot_global_index_text = ast_runtime_for_static_slot_global_index_probe.compile().mlir_text() expect_parse_roundtrip_and_verify( ast_runtime_for_static_slot_global_index_text,