fix(ptodsl): correct With slot liveness

ptodsl/ptodsl/_ast_rewrite.py

@@ -601,6 +601,16 @@ 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)):
+        context_loads = set()
+        bound_bases = set()
+        for item in stmt.items:
+            context_loads |= _slot_info(item.context_expr, static_env, static_iters).loads
+            if item.optional_vars is not None:
+                bound_bases |= _simple_name_targets(item.optional_vars)
+        body_live = _slot_live_before_block(stmt.body, live_after, static_env, static_iters)
+        body_live = {slot for slot in body_live if slot.base not in bound_bases}
+        return context_loads | body_live
     if isinstance(stmt, ast.If):
         test_info = _slot_info(stmt.test, static_env, static_iters)
         return (

ptodsl/tests/test_jit_compile.py

@@ -2117,6 +2117,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 +7123,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,