diff --git a/python/flydsl/compiler/ast_rewriter.py b/python/flydsl/compiler/ast_rewriter.py index dddd24b9b..1eff34554 100644 --- a/python/flydsl/compiler/ast_rewriter.py +++ b/python/flydsl/compiler/ast_rewriter.py @@ -18,6 +18,7 @@ from ..expr.meta import capture_user_location, file_location from ..expr.typing import as_dsl_value, as_ir_value from ..utils import env, log +from .protocol import construct_from_ir_values def _set_lineno(node, n=1): @@ -58,6 +59,167 @@ def _unwrap_constexpr(node): return node +# --------------------------------------------------------------------------- +# Explode / assemble Python containers for control-flow carried state. +# +# A dynamic scf.if / scf.for / scf.while can only carry a *flat* list of +# ir.Value as its results / iter_args / block_args. A carried Python variable, +# however, may be a *container* (list / tuple / dict / SimpleNamespace, possibly +# nested) whose elements are DSL values. These helpers explode such a container +# into the flat ir.Value list the scf op needs, and assemble the original +# container shape back on the way out. +# +# The structural template used for assembly is the *pre-region value* itself: +# it already records the container nesting and each element's dtype. Elements +# reuse the existing DSL<->ir.Value protocol (as_ir_value / as_dsl_value / +# protocol.construct_from_ir_values), so scalars, Vectors and @fx.struct values +# all keep working; a bare scalar is just the degenerate single-slot case, so +# previous (scalar-only) behaviour is preserved. +# --------------------------------------------------------------------------- + + +def _carried_slot_count(element): + """Number of ir.Value slots a single (non-container) carried element occupies. + + Pure: it must never materialize new IR (so it is safe to call for slicing). + """ + if isinstance(element, ir.Value): + return 1 + if hasattr(element, "__get_ir_types__"): + return len(element.__get_ir_types__()) + if hasattr(element, "__extract_to_ir_values__"): + # extract returns the element's *existing* ir.Values, no new ops are built + return len(element.__extract_to_ir_values__()) + # bare python scalar -> as_ir_value would coerce it to a single constant + return 1 + + +def _explode_carried(value, out): + """Depth-first explode ``value`` into ``out`` (a list of ir.Value slots). + + Containers (list/tuple/dict/SimpleNamespace) are recursed; every other + object is treated as a single element and lowered with ``as_ir_value`` so + that the exact per-element coercion (e.g. python int -> constant, multi-value + struct -> several ir.Values) matches the rest of the codebase. + """ + if isinstance(value, (list, tuple)): + for v in value: + _explode_carried(v, out) + elif isinstance(value, dict): + for v in value.values(): + _explode_carried(v, out) + elif isinstance(value, types.SimpleNamespace): + for v in vars(value).values(): + _explode_carried(v, out) + else: + raw = as_ir_value(value) + # a multi-value element (e.g. an @fx.struct) extracts to a python list + if isinstance(raw, list): + out.extend(raw) + else: + out.append(raw) + + +def _assemble_carried(exemplar, slots, cursor): + """Inverse of :func:`_explode_carried`. + + Assemble the container shape of ``exemplar`` consuming ``slots[cursor:]``; + returns ``(value, next_cursor)``. ``exemplar`` supplies both the nesting and, + at the elements, the DSL type to wrap each ir.Value slot back into. + """ + if isinstance(exemplar, (list, tuple)): + items = [] + for ex in exemplar: + v, cursor = _assemble_carried(ex, slots, cursor) + items.append(v) + return type(exemplar)(items), cursor + if isinstance(exemplar, dict): + built = {} + for k, ex in exemplar.items(): + built[k], cursor = _assemble_carried(ex, slots, cursor) + return built, cursor + if isinstance(exemplar, types.SimpleNamespace): + built = {} + for k, ex in vars(exemplar).items(): + built[k], cursor = _assemble_carried(ex, slots, cursor) + return types.SimpleNamespace(**built), cursor + # single element + n = _carried_slot_count(exemplar) + seg = list(slots[cursor : cursor + n]) + cursor += n + if n == 1: + # scalar element: identical to the previous per-name packing + element = as_dsl_value(seg[0], exemplar) + else: + # multi-value element (e.g. @fx.struct): rebuild via the DSL protocol + element = construct_from_ir_values(type(exemplar), exemplar, seg) + return element, cursor + + +def _explode_states(names, values, ctx_label): + """Explode every carried ``value`` into a shared flat ir.Value slot list. + + Returns ``(state_raw, counts, exemplars, result_types)`` where + ``counts[i]`` is how many flat ir.Value slots ``names[i]`` contributed and + ``exemplars[i]`` is its pre-region value (the assembly template). + """ + state_raw = [] + counts = [] + exemplars = [] + for name, value in zip(names, values): + slots = [] + _explode_carried(value, slots) + for slot in slots: + if not isinstance(slot, ir.Value): + raise TypeError( + f"{ctx_label} variable '{name}' contains a {type(slot).__name__} " + "that is not an MLIR Value; only MLIR-backed elements can be carried " + "through dynamic control flow. Initialize the element with a DSL value " + "(e.g. fx.Int32(0), fx.Float32(0.0))." + ) + counts.append(len(slots)) + exemplars.append(value) + state_raw.extend(slots) + result_types = [v.type for v in state_raw] + return state_raw, counts, exemplars, result_types + + +def _assemble_states(slots, counts, exemplars): + """Split ``slots`` per carried name (by ``counts``) and rebuild each value + from its ``exemplar`` template. Inverse of :func:`_explode_states`.""" + out = [] + cursor = 0 + for count, exemplar in zip(counts, exemplars): + seg = list(slots[cursor : cursor + count]) + cursor += count + value, _ = _assemble_carried(exemplar, seg, 0) + out.append(value) + return out + + +def _explode_branch_outputs(names, values, counts, result_types, ctx_label): + """Explode a branch/body's produced values and verify they match the region + boundary exactly (same slot count and same per-slot dtype), so scf.if/for/ + while operand/result unification stays legal and user mistakes fail early.""" + out_raw, out_counts, _exemplars, out_types = _explode_states(names, values, ctx_label) + for name, c_in, c_out in zip(names, counts, out_counts): + if c_in != c_out: + raise TypeError( + f"{ctx_label} variable '{name}' changed container shape across the " + f"region boundary: entered with {c_in} MLIR value slot(s) but the " + f"branch/body produced {c_out}. A carried container's length/nesting " + "must be identical on entry and exit." + ) + for slot, (t_in, t_out) in enumerate(zip(result_types, out_types)): + if t_in != t_out: + raise TypeError( + f"{ctx_label}: carried element type mismatch at slot {slot}: " + f"expected {t_in}, got {t_out}. A carried element's dtype must match on " + "entry and exit." + ) + return out_raw + + def _collect_assigned_vars(body_stmts, active_symbols, orelse_stmts=None, test_expr=None): write_args = [] invoked_args = [] @@ -659,54 +821,28 @@ def scf_if_dispatch( if else_fn is None: else_fn = lambda *args: {} - state_raw = [] - for name, value in zip(result_names, result_values): - raw = as_ir_value(value) - if not isinstance(raw, ir.Value): - raise TypeError( - f"state variable '{name}' is {type(raw).__name__}, not an MLIR Value; " - "stateful dynamic if requires MLIR-backed values." - ) - state_raw.append(raw) - - result_types = [v.type for v in state_raw] + # Explode each carried value (possibly a list/tuple/dict/nested + # container) into the flat ir.Value slots that scf.if can carry as + # results; `counts`/`exemplars` let us assemble the containers on exit. + state_raw, counts, exemplars, result_types = _explode_states(result_names, result_values, "dynamic if/else") if_op = scf.IfOp(cond_i1, result_types, has_else=True, loc=capture_user_location()) - with ir.InsertionPoint(if_op.regions[0].blocks[0]): - then_result = ReplaceIfWithDispatch._call_branch(then_fn, result_names, result_values) - then_values = ReplaceIfWithDispatch._normalize_branch_result( - then_result, result_names, result_map, "then-branch" - ) - then_raw = ReplaceIfWithDispatch._unwrap_mlir_values(then_values, result_names, "then-branch") - for name, expect_ty, got in zip(result_names, result_types, then_raw): - if got.type != expect_ty: - raise TypeError( - f"if/else variable '{name}' type mismatch in then-branch: " - f"expected {expect_ty}, got {got.type}" - ) - scf.YieldOp(then_raw, loc=capture_user_location()) + def _emit_branch(fn, block, label): + with ir.InsertionPoint(block): + branch_result = ReplaceIfWithDispatch._call_branch(fn, result_names, result_values) + branch_values = ReplaceIfWithDispatch._normalize_branch_result( + branch_result, result_names, result_map, label + ) + branch_raw = _explode_branch_outputs(result_names, branch_values, counts, result_types, label) + scf.YieldOp(branch_raw, loc=capture_user_location()) + _emit_branch(then_fn, if_op.regions[0].blocks[0], "then-branch") if len(if_op.regions[1].blocks) == 0: if_op.regions[1].blocks.append(*[]) - with ir.InsertionPoint(if_op.regions[1].blocks[0]): - else_result = ReplaceIfWithDispatch._call_branch(else_fn, result_names, result_values) - else_values = ReplaceIfWithDispatch._normalize_branch_result( - else_result, result_names, result_map, "else-branch" - ) - else_raw = ReplaceIfWithDispatch._unwrap_mlir_values(else_values, result_names, "else-branch") - for name, expect_ty, got in zip(result_names, result_types, else_raw): - if got.type != expect_ty: - raise TypeError( - f"if/else variable '{name}' type mismatch in else-branch: " - f"expected {expect_ty}, got {got.type}" - ) - scf.YieldOp(else_raw, loc=capture_user_location()) + _emit_branch(else_fn, if_op.regions[1].blocks[0], "else-branch") - wrapped = ReplaceIfWithDispatch._pack_dispatch_results(list(if_op.results), result_values) - if len(result_names) == 1: - final_values = [wrapped] - else: - final_values = list(wrapped) + # Re-pack the flat scf.if results back into each name's container shape. + final_values = _assemble_states(list(if_op.results), counts, exemplars) return ReplaceIfWithDispatch._pack_named_values(result_names, final_values) @classmethod @@ -886,39 +1022,34 @@ def scf_ifexp_dispatch(cond, then_fn, else_fn): if not isinstance(cond_i1, ir.Value): raise TypeError(f"dynamic ifexp condition must lower to ir.Value, got {type(cond_i1).__name__}") + # A ternary may evaluate to a *container* (e.g. `[a, b] if c else [x, y]`), + # so treat its value like a single carried state entry and explode it. + # Probe both branches in a throwaway region to learn the flat structure + # (slot count + dtypes) and the container template for assembly. + NAME = ("",) sandbox = scf.ExecuteRegionOp(result=[]) sandbox.region.blocks.append() with ir.InsertionPoint(sandbox.region.blocks[0]): probe_then = then_fn() - probe_then_raw = as_ir_value(probe_then) - probe_else = else_fn() - probe_else_raw = as_ir_value(probe_else) - if not isinstance(probe_then_raw, ir.Value): - raise TypeError( - f"dynamic ifexp then-branch must produce an MLIR Value, " f"got {type(probe_then_raw).__name__}" - ) - if not isinstance(probe_else_raw, ir.Value): - raise TypeError( - f"dynamic ifexp else-branch must produce an MLIR Value, " f"got {type(probe_else_raw).__name__}" - ) - if probe_then_raw.type != probe_else_raw.type: - raise TypeError( - f"dynamic ifexp type mismatch: " - f"then-branch produces {probe_then_raw.type}, " - f"else-branch produces {probe_else_raw.type}" - ) - yield_type = probe_then_raw.type + _then_raw, counts, exemplars, yield_types = _explode_states( + NAME, (probe_then,), "dynamic ifexp then-branch" + ) + # else must produce the same structure and element dtypes as then + _explode_branch_outputs(NAME, (else_fn(),), counts, yield_types, "dynamic ifexp else-branch") + sandbox.operation.erase() - op = scf.IfOp(cond_i1, [yield_type], has_else=True, loc=capture_user_location()) + op = scf.IfOp(cond_i1, yield_types, has_else=True, loc=capture_user_location()) with ir.InsertionPoint(op.regions[0].blocks[0]): - scf.YieldOp([as_ir_value(then_fn())], loc=capture_user_location()) + then_raw = _explode_branch_outputs(NAME, (then_fn(),), counts, yield_types, "dynamic ifexp then-branch") + scf.YieldOp(then_raw, loc=capture_user_location()) if len(op.regions[1].blocks) == 0: op.regions[1].blocks.append() with ir.InsertionPoint(op.regions[1].blocks[0]): - scf.YieldOp([as_ir_value(else_fn())], loc=capture_user_location()) + else_raw = _explode_branch_outputs(NAME, (else_fn(),), counts, yield_types, "dynamic ifexp else-branch") + scf.YieldOp(else_raw, loc=capture_user_location()) - sandbox.operation.erase() - return as_dsl_value(op.results[0], probe_then) + # Assemble the flat scf.if results into the ternary's (possibly container) value. + return _assemble_states(list(op.results), counts, exemplars)[0] @ASTRewriter.register @@ -1008,15 +1139,9 @@ def scf_for_dispatch(start, stop, step, body_fn, *, result_names=(), result_valu scf.YieldOp([], loc=capture_user_location()) return ReplaceIfWithDispatch._pack_named_values(result_names, result_values) - state_raw = [] - for name, value in zip(result_names, result_values): - raw = as_ir_value(value) - if not isinstance(raw, ir.Value): - raise TypeError( - f"for-loop variable '{name}' is {type(raw).__name__}, not an MLIR Value; " - "only MLIR-backed values can be loop-carried in dynamic for loops." - ) - state_raw.append(raw) + # Explode carried values (possibly containers) into the flat iter_args + # scf.for needs; assemble containers for the body and for the results. + state_raw, counts, exemplars, result_types = _explode_states(result_names, result_values, "dynamic for loop") loc = capture_user_location() for_op = scf.ForOp(start_val, stop_val, step_val, state_raw, loc=loc) @@ -1024,27 +1149,19 @@ def scf_for_dispatch(start, stop, step, body_fn, *, result_names=(), result_valu with ir.InsertionPoint(for_op.body): iv = for_op.induction_variable - inner_args = [as_dsl_value(a, ex) for a, ex in zip(for_op.inner_iter_args, result_values)] + # inner_iter_args are flat; assemble them into each name's container + # so the body sees the same list/tuple/... shape it wrote. + inner_args = _assemble_states(list(for_op.inner_iter_args), counts, exemplars) body_result = body_fn(iv, result_names, *inner_args) body_values = ReplaceIfWithDispatch._normalize_branch_result( body_result, result_names, result_map, "for-body" ) - body_raw = ReplaceIfWithDispatch._unwrap_mlir_values(body_values, result_names, "for-body") - result_types = [v.type for v in state_raw] - for name, expect_ty, got in zip(result_names, result_types, body_raw): - if got.type != expect_ty: - raise TypeError( - f"for-loop variable '{name}' type mismatch: " f"expected {expect_ty}, got {got.type}" - ) + body_raw = _explode_branch_outputs(result_names, body_values, counts, result_types, "for-body") scf.YieldOp(body_raw, loc=capture_user_location()) - wrapped = ReplaceIfWithDispatch._pack_dispatch_results(list(for_op.results), result_values) - if len(result_names) == 1: - final_values = [wrapped] - else: - final_values = list(wrapped) + final_values = _assemble_states(list(for_op.results), counts, exemplars) return ReplaceIfWithDispatch._pack_named_values(result_names, final_values) @classmethod @@ -1334,17 +1451,9 @@ def scf_while_dispatch(before_fn, after_fn, *, result_names=(), result_values=() f"const_expr() if it is a compile-time constant." ) - state_raw = [] - for name, value in zip(result_names, result_values): - raw = as_ir_value(value) - if not isinstance(raw, ir.Value): - raise TypeError( - f"while-loop variable '{name}' is {type(raw).__name__}, not an MLIR Value; " - "only MLIR-backed values can be loop-carried in dynamic while loops." - ) - state_raw.append(raw) - - result_types = [v.type for v in state_raw] + # Explode carried values (possibly containers) into the flat state the + # scf.while before/after blocks carry; assemble containers per block. + state_raw, counts, exemplars, result_types = _explode_states(result_names, result_values, "dynamic while loop") loc = capture_user_location() while_op = scf.WhileOp(result_types, state_raw, loc=loc) # Give the loop-carried block arguments the user location. @@ -1354,7 +1463,7 @@ def scf_while_dispatch(before_fn, after_fn, *, result_names=(), result_values=() with ir.InsertionPoint(while_op.regions[0].blocks[0]): before_args = list(while_op.regions[0].blocks[0].arguments) - wrapped_before = [as_dsl_value(a, ex) for a, ex in zip(before_args, result_values)] if result_names else [] + wrapped_before = _assemble_states(before_args, counts, exemplars) if result_names else [] before_cond = ReplaceIfWithDispatch._call_branch(before_fn, result_names, wrapped_before) cond_i1 = ReplaceIfWithDispatch._to_i1(before_cond) if not isinstance(cond_i1, ir.Value): @@ -1363,18 +1472,13 @@ def scf_while_dispatch(before_fn, after_fn, *, result_names=(), result_values=() with ir.InsertionPoint(while_op.regions[1].blocks[0]): after_args = list(while_op.regions[1].blocks[0].arguments) - wrapped_after = [as_dsl_value(a, ex) for a, ex in zip(after_args, result_values)] if result_names else [] + wrapped_after = _assemble_states(after_args, counts, exemplars) if result_names else [] body_result = ReplaceIfWithDispatch._call_branch(after_fn, result_names, wrapped_after) if result_names: body_values = ReplaceIfWithDispatch._normalize_branch_result( body_result, result_names, result_map, "while-body" ) - body_raw = ReplaceIfWithDispatch._unwrap_mlir_values(body_values, result_names, "while-body") - for name, expect_ty, got in zip(result_names, result_types, body_raw): - if got.type != expect_ty: - raise TypeError( - f"while-loop variable '{name}' type mismatch: expected {expect_ty}, got {got.type}" - ) + body_raw = _explode_branch_outputs(result_names, body_values, counts, result_types, "while-body") scf.YieldOp(body_raw, loc=capture_user_location()) else: scf.YieldOp([], loc=capture_user_location()) @@ -1382,11 +1486,7 @@ def scf_while_dispatch(before_fn, after_fn, *, result_names=(), result_values=() if not result_names: return ReplaceIfWithDispatch._pack_named_values(result_names, result_values) - wrapped = ReplaceIfWithDispatch._pack_dispatch_results(list(while_op.results), result_values) - if len(result_names) == 1: - final_values = [wrapped] - else: - final_values = list(wrapped) + final_values = _assemble_states(list(while_op.results), counts, exemplars) return ReplaceIfWithDispatch._pack_named_values(result_names, final_values) @classmethod diff --git a/tests/system/test_dynamic_controlflow_list_carry_e2e.py b/tests/system/test_dynamic_controlflow_list_carry_e2e.py new file mode 100644 index 000000000..8c5e5206e --- /dev/null +++ b/tests/system/test_dynamic_controlflow_list_carry_e2e.py @@ -0,0 +1,141 @@ +#!/usr/bin/env python3 + +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2025 FlyDSL Project Contributors + +"""End-to-end tests: carry a Python list through dynamic if/for/while/ifexp in a +real @flyc.kernel, launch it, and check the results against the expected values. + +This is the end-to-end counterpart of the MLIR-level unit tests in +tests/unit/test_dynamic_controlflow_list_carry.py. The pattern (a list local +reassigned inside a runtime-conditioned region) is the one that previously +failed to compile with ``TypeError: ... is list, not an MLIR Value``. +""" + +import pytest +import torch + +import flydsl.compiler as flyc +import flydsl.expr as fx + +pytestmark = [pytest.mark.l2_device, pytest.mark.rocm_lower] + +if not torch.cuda.is_available(): + pytest.skip("CUDA/ROCm not available", allow_module_level=True) + + +def _out(n): + out = torch.zeros(n, device="cuda", dtype=torch.int32) + t_out = flyc.from_torch_tensor(out).mark_layout_dynamic(leading_dim=0, divisibility=1) + return out, t_out + + +# ── dynamic if carrying a list ────────────────────────────────────────────── + + +@flyc.kernel +def _kernel_if_list(Out: fx.Tensor, flag: fx.Int32): + lst = [fx.Int32(1), fx.Int32(2)] + if flag > fx.Int32(0): # runtime condition -> dynamic scf.if + lst = [lst[0] + fx.Int32(10), lst[1] + fx.Int32(20)] + rsrc = fx.buffer_ops.create_buffer_resource(Out) + fx.buffer_ops.buffer_store(lst[0], rsrc, fx.Int32(0)) + fx.buffer_ops.buffer_store(lst[1], rsrc, fx.Int32(1)) + + +@flyc.jit +def _run_if_list(Out: fx.Tensor, flag: fx.Int32, stream: fx.Stream = fx.Stream(None)): + _kernel_if_list(Out, flag).launch(grid=(1, 1, 1), block=(1, 1, 1), stream=stream.value) + + +# ── dynamic for carrying a list ───────────────────────────────────────────── + + +@flyc.kernel +def _kernel_for_list(Out: fx.Tensor, n: fx.Int32): + lst = [fx.Int32(0), fx.Int32(100)] + for i in range(n): + lst = [lst[0] + fx.Int32(1), lst[1] - fx.Int32(1)] + rsrc = fx.buffer_ops.create_buffer_resource(Out) + fx.buffer_ops.buffer_store(lst[0], rsrc, fx.Int32(0)) + fx.buffer_ops.buffer_store(lst[1], rsrc, fx.Int32(1)) + + +@flyc.jit +def _run_for_list(Out: fx.Tensor, n: fx.Int32, stream: fx.Stream = fx.Stream(None)): + _kernel_for_list(Out, n).launch(grid=(1, 1, 1), block=(1, 1, 1), stream=stream.value) + + +# ── dynamic while carrying a list ─────────────────────────────────────────── + + +@flyc.kernel +def _kernel_while_list(Out: fx.Tensor, n: fx.Int32): + lst = [n, fx.Int32(0)] + while lst[0] > fx.Int32(0): + lst = [lst[0] - fx.Int32(1), lst[1] + fx.Int32(1)] + rsrc = fx.buffer_ops.create_buffer_resource(Out) + fx.buffer_ops.buffer_store(lst[0], rsrc, fx.Int32(0)) + fx.buffer_ops.buffer_store(lst[1], rsrc, fx.Int32(1)) + + +@flyc.jit +def _run_while_list(Out: fx.Tensor, n: fx.Int32, stream: fx.Stream = fx.Stream(None)): + _kernel_while_list(Out, n).launch(grid=(1, 1, 1), block=(1, 1, 1), stream=stream.value) + + +# ── dynamic ifexp (ternary) evaluating to a list ──────────────────────────── + + +@flyc.kernel +def _kernel_ifexp_list(Out: fx.Tensor, flag: fx.Int32): + lst = [fx.Int32(1), fx.Int32(2)] if flag > fx.Int32(0) else [fx.Int32(3), fx.Int32(4)] + rsrc = fx.buffer_ops.create_buffer_resource(Out) + fx.buffer_ops.buffer_store(lst[0], rsrc, fx.Int32(0)) + fx.buffer_ops.buffer_store(lst[1], rsrc, fx.Int32(1)) + + +@flyc.jit +def _run_ifexp_list(Out: fx.Tensor, flag: fx.Int32, stream: fx.Stream = fx.Stream(None)): + _kernel_ifexp_list(Out, flag).launch(grid=(1, 1, 1), block=(1, 1, 1), stream=stream.value) + + +# ── Tests ─────────────────────────────────────────────────────────────────── + + +class TestDynamicControlFlowListCarryE2E: + def test_if_list_taken(self): + out, t_out = _out(2) + _run_if_list(t_out, fx.Int32(1)) # flag>0 -> then branch + torch.cuda.synchronize() + assert out[0].item() == 11 and out[1].item() == 22, out.tolist() + + def test_if_list_not_taken(self): + out, t_out = _out(2) + _run_if_list(t_out, fx.Int32(0)) # flag==0 -> else keeps original list + torch.cuda.synchronize() + assert out[0].item() == 1 and out[1].item() == 2, out.tolist() + + def test_for_list(self): + out, t_out = _out(2) + _run_for_list(t_out, fx.Int32(5)) # +1/-1 x5 -> [5, 95] + torch.cuda.synchronize() + assert out[0].item() == 5 and out[1].item() == 95, out.tolist() + + def test_while_list(self): + out, t_out = _out(2) + _run_while_list(t_out, fx.Int32(7)) # drain 7 -> [0, 7] + torch.cuda.synchronize() + assert out[0].item() == 0 and out[1].item() == 7, out.tolist() + + def test_ifexp_list_taken(self): + out, t_out = _out(2) + _run_ifexp_list(t_out, fx.Int32(1)) # flag>0 -> [1, 2] + torch.cuda.synchronize() + assert out[0].item() == 1 and out[1].item() == 2, out.tolist() + + def test_ifexp_list_not_taken(self): + out, t_out = _out(2) + _run_ifexp_list(t_out, fx.Int32(0)) # else -> [3, 4] + torch.cuda.synchronize() + assert out[0].item() == 3 and out[1].item() == 4, out.tolist() diff --git a/tests/unit/test_dynamic_controlflow_list_carry.py b/tests/unit/test_dynamic_controlflow_list_carry.py new file mode 100644 index 000000000..c21f70dd0 --- /dev/null +++ b/tests/unit/test_dynamic_controlflow_list_carry.py @@ -0,0 +1,342 @@ +#!/usr/bin/env python3 + +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2025 FlyDSL Project Contributors + +"""MLIR-level unit tests (no GPU) for carrying Python containers (list / tuple / +dict / nested) through dynamic scf.if / scf.for / scf.while. + +Before this feature a reassigned container local tripped the dispatchers with +``TypeError: ... is list, not an MLIR Value``; the region rewriter could only +carry a single ir.Value per name. These tests lock in the container round-trip +(explode to per-element iter_args/results, assemble on exit) and the guard rails +when a branch/body changes a container's shape or an element's dtype. +""" + +import pytest + +from flydsl._mlir.dialects import arith, func +from flydsl._mlir.ir import Context, FunctionType, InsertionPoint, IntegerType, Location, Module +from flydsl.compiler.ast_rewriter import ( + CanonicalizeWhile, + InsertEmptyYieldForSCFFor, + ReplaceIfWithDispatch, +) +from flydsl.expr.numeric import Int32, Int64 + + +def _i32(v): + return Int32(arith.ConstantOp(IntegerType.get_signless(32), v).result) + + +def _i64(v): + return Int64(arith.ConstantOp(IntegerType.get_signless(64), v).result) + + +# ─────────────────────────────── if ──────────────────────────────────────── + + +def test_if_carries_list(): + """A reassigned list local carried across a dynamic scf.if. Result must be a + Python list of the same length, and the scf.if must carry one result per list + element.""" + with Context(), Location.unknown(): + module = Module.create() + i1 = IntegerType.get_signless(1) + with InsertionPoint(module.body): + f = func.FuncOp("if_list", FunctionType.get([i1], [])) + entry = f.add_entry_block() + with InsertionPoint(entry): + cond = entry.arguments[0] + lst = [_i32(0), _i32(1)] + + out = ReplaceIfWithDispatch.scf_if_dispatch( + cond, + lambda names, lst: {"lst": [_i32(10), _i32(11)]}, + lambda names, lst: {"lst": [_i32(20), _i32(21)]}, + result_names=("lst",), + result_values=(lst,), + ) + assert isinstance(out, list) and len(out) == 2 + assert all(isinstance(x, Int32) for x in out) + func.ReturnOp([]) + + text = str(module) + assert module.operation.verify() + assert "scf.if" in text + # one scf.if carrying two i32 results (one per list element) + assert "-> (i32, i32)" in text + + +def test_if_carries_nested_list(): + with Context(), Location.unknown(): + module = Module.create() + i1 = IntegerType.get_signless(1) + with InsertionPoint(module.body): + f = func.FuncOp("if_nested", FunctionType.get([i1], [])) + entry = f.add_entry_block() + with InsertionPoint(entry): + cond = entry.arguments[0] + nested = [[_i32(0), _i32(1)], [_i32(2)]] # 3 elements + + out = ReplaceIfWithDispatch.scf_if_dispatch( + cond, + lambda names, nested: {"nested": [[_i32(10), _i32(11)], [_i32(12)]]}, + lambda names, nested: {"nested": [[_i32(20), _i32(21)], [_i32(22)]]}, + result_names=("nested",), + result_values=(nested,), + ) + assert isinstance(out, list) and len(out) == 2 + assert len(out[0]) == 2 and len(out[1]) == 1 + func.ReturnOp([]) + + assert module.operation.verify() + assert "-> (i32, i32, i32)" in str(module) + + +def test_if_carries_tuple_and_dict(): + with Context(), Location.unknown(): + module = Module.create() + i1 = IntegerType.get_signless(1) + with InsertionPoint(module.body): + f = func.FuncOp("if_tuple_dict", FunctionType.get([i1], [])) + entry = f.add_entry_block() + with InsertionPoint(entry): + cond = entry.arguments[0] + tup = (_i32(0), _i32(1)) + dct = {"x": _i32(2), "y": _i32(3)} + + out = ReplaceIfWithDispatch.scf_if_dispatch( + cond, + lambda names, tup, dct: {"tup": (_i32(10), _i32(11)), "dct": {"x": _i32(12), "y": _i32(13)}}, + lambda names, tup, dct: {"tup": (_i32(20), _i32(21)), "dct": {"x": _i32(22), "y": _i32(23)}}, + result_names=("tup", "dct"), + result_values=(tup, dct), + ) + rtup, rdct = out + assert isinstance(rtup, tuple) and len(rtup) == 2 + assert isinstance(rdct, dict) and set(rdct) == {"x", "y"} + func.ReturnOp([]) + + assert module.operation.verify() + + +def test_if_mixes_list_and_scalar(): + """A list carry and a plain scalar carry in the same if must both work + (scalar is the degenerate single-element case).""" + with Context(), Location.unknown(): + module = Module.create() + i1 = IntegerType.get_signless(1) + with InsertionPoint(module.body): + f = func.FuncOp("if_mixed", FunctionType.get([i1], [])) + entry = f.add_entry_block() + with InsertionPoint(entry): + cond = entry.arguments[0] + lst = [_i32(0), _i32(1)] + s = _i32(7) + + out = ReplaceIfWithDispatch.scf_if_dispatch( + cond, + lambda names, lst, s: {"lst": [_i32(10), _i32(11)], "s": _i32(70)}, + lambda names, lst, s: {"lst": [_i32(20), _i32(21)], "s": _i32(80)}, + result_names=("lst", "s"), + result_values=(lst, s), + ) + rlst, rs = out + assert isinstance(rlst, list) and len(rlst) == 2 + assert isinstance(rs, Int32) + func.ReturnOp([]) + + assert module.operation.verify() + assert "-> (i32, i32, i32)" in str(module) + + +def test_if_static_cond_with_list_no_ifop(): + """A const_expr (python-bool) condition must resolve at trace time: pick a + branch, build no scf.if, still return a list.""" + with Context(), Location.unknown(): + module = Module.create() + with InsertionPoint(module.body): + f = func.FuncOp("if_static_list", FunctionType.get([], [])) + entry = f.add_entry_block() + with InsertionPoint(entry): + lst = [_i32(0), _i32(1)] + out = ReplaceIfWithDispatch.scf_if_dispatch( + True, + lambda names, lst: {"lst": [_i32(10), _i32(11)]}, + lambda names, lst: {"lst": [_i32(20), _i32(21)]}, + result_names=("lst",), + result_values=(lst,), + ) + assert isinstance(out, list) and len(out) == 2 + func.ReturnOp([]) + + assert module.operation.verify() + assert "scf.if" not in str(module) + + +def test_if_list_length_change_errors(): + """A branch that changes the carried list's length must fail with a clear + error (scf.if result arity would otherwise mismatch).""" + with Context(), Location.unknown(): + module = Module.create() + i1 = IntegerType.get_signless(1) + with InsertionPoint(module.body): + f = func.FuncOp("if_badlen", FunctionType.get([i1], [])) + entry = f.add_entry_block() + with InsertionPoint(entry): + cond = entry.arguments[0] + lst = [_i32(0), _i32(1)] + with pytest.raises(TypeError, match="changed container shape"): + ReplaceIfWithDispatch.scf_if_dispatch( + cond, + lambda names, lst: {"lst": [_i32(10)]}, # length 1 != 2 + lambda names, lst: {"lst": [_i32(20), _i32(21)]}, + result_names=("lst",), + result_values=(lst,), + ) + + +def test_if_list_dtype_change_errors(): + """A branch that changes an element's dtype must fail with a clear error.""" + with Context(), Location.unknown(): + module = Module.create() + i1 = IntegerType.get_signless(1) + with InsertionPoint(module.body): + f = func.FuncOp("if_baddtype", FunctionType.get([i1], [])) + entry = f.add_entry_block() + with InsertionPoint(entry): + cond = entry.arguments[0] + lst = [_i32(0)] + with pytest.raises(TypeError, match="type mismatch"): + ReplaceIfWithDispatch.scf_if_dispatch( + cond, + lambda names, lst: {"lst": [_i64(10)]}, # i64 != i32 + lambda names, lst: {"lst": [_i32(20)]}, + result_names=("lst",), + result_values=(lst,), + ) + + +# ─────────────────────────────── ifexp ───────────────────────────────────── + + +def test_ifexp_carries_list(): + """A ternary whose value is a list: `[a, b] if cond else [c, d]`.""" + with Context(), Location.unknown(): + module = Module.create() + i1 = IntegerType.get_signless(1) + with InsertionPoint(module.body): + f = func.FuncOp("ifexp_list", FunctionType.get([i1], [])) + entry = f.add_entry_block() + with InsertionPoint(entry): + cond = entry.arguments[0] + out = ReplaceIfWithDispatch.scf_ifexp_dispatch( + cond, lambda: [_i32(1), _i32(2)], lambda: [_i32(3), _i32(4)] + ) + assert isinstance(out, list) and len(out) == 2 + assert all(isinstance(x, Int32) for x in out) + func.ReturnOp([]) + + text = str(module) + assert module.operation.verify() + assert "scf.if" in text + assert "-> (i32, i32)" in text + + +def test_ifexp_scalar_still_scalar(): + """Scalar ternary keeps returning a scalar (backward compatible).""" + with Context(), Location.unknown(): + module = Module.create() + i32 = IntegerType.get_signless(32) + i1 = IntegerType.get_signless(1) + with InsertionPoint(module.body): + f = func.FuncOp("ifexp_scalar", FunctionType.get([i1], [i32])) + entry = f.add_entry_block() + with InsertionPoint(entry): + cond = entry.arguments[0] + out = ReplaceIfWithDispatch.scf_ifexp_dispatch(cond, lambda: _i32(1), lambda: _i32(2)) + assert isinstance(out, Int32) + func.ReturnOp([out.ir_value()]) + + assert module.operation.verify() + assert "-> (i32)" in str(module) + + +def test_ifexp_shape_mismatch_errors(): + """then/else producing different container shapes must fail clearly.""" + with Context(), Location.unknown(): + module = Module.create() + i1 = IntegerType.get_signless(1) + with InsertionPoint(module.body): + f = func.FuncOp("ifexp_bad", FunctionType.get([i1], [])) + entry = f.add_entry_block() + with InsertionPoint(entry): + cond = entry.arguments[0] + with pytest.raises(TypeError, match="changed container shape"): + ReplaceIfWithDispatch.scf_ifexp_dispatch( + cond, lambda: [_i32(1), _i32(2)], lambda: [_i32(3)] # len 1 != 2 + ) + + +# ─────────────────────────────── for ─────────────────────────────────────── + + +def test_for_carries_list(): + """for i in range(4): lst = [lst[0]+1, lst[1]+2] → list survives as iter_args.""" + with Context(), Location.unknown(): + module = Module.create() + i32 = IntegerType.get_signless(32) + with InsertionPoint(module.body): + f = func.FuncOp("for_list", FunctionType.get([], [i32, i32])) + entry = f.add_entry_block() + with InsertionPoint(entry): + lst = [_i32(0), _i32(1)] + + def body_fn(iv, names, lst): + # body receives the reassembled list + assert isinstance(lst, list) and len(lst) == 2 + return {"lst": [lst[0] + _i32(1), lst[1] + _i32(2)]} + + out = InsertEmptyYieldForSCFFor.scf_for_dispatch( + 0, 4, 1, body_fn, result_names=("lst",), result_values=(lst,) + ) + assert isinstance(out, list) and len(out) == 2 + func.ReturnOp([out[0].ir_value(), out[1].ir_value()]) + + text = str(module) + assert module.operation.verify() + assert "scf.for" in text + assert "-> (i32, i32)" in text + + +# ─────────────────────────────── while ───────────────────────────────────── + + +def test_while_carries_list(): + """while lst[0] > 0: lst = [lst[0]-1, lst[1]+1] → list survives before/after.""" + with Context(), Location.unknown(): + module = Module.create() + i32 = IntegerType.get_signless(32) + with InsertionPoint(module.body): + f = func.FuncOp("while_list", FunctionType.get([], [i32, i32])) + entry = f.add_entry_block() + with InsertionPoint(entry): + lst = [_i32(4), _i32(0)] + + def before_fn(names, lst): + assert isinstance(lst, list) and len(lst) == 2 + return lst[0] > _i32(0) + + def after_fn(names, lst): + return {"lst": [lst[0] - _i32(1), lst[1] + _i32(1)]} + + out = CanonicalizeWhile.scf_while_dispatch( + before_fn, after_fn, result_names=("lst",), result_values=(lst,) + ) + assert isinstance(out, list) and len(out) == 2 + func.ReturnOp([out[0].ir_value(), out[1].ir_value()]) + + text = str(module) + assert module.operation.verify() + assert "scf.while" in text