From 8630a951e9ee5360ac89a044b8e8f32e79fc2133 Mon Sep 17 00:00:00 2001 From: xudoyuan Date: Mon, 20 Jul 2026 10:23:59 +0000 Subject: [PATCH 1/4] [compiler] carry Python containers through dynamic if/for/while/ifexp Dynamic scf.if / scf.for / scf.while / ifexp previously required every carried variable to be a single ir.Value, so a reassigned list/tuple/dict local raised "state variable '...' is list, not an MLIR Value". Route the four dispatchers through an explode/assemble step that reuses the existing DSL<->ir.Value protocol (protocol.extract/construct): each carried value (possibly a nested container) is exploded to per-element iter_args/ results and assembled back from the pre-region value used as the structural template on exit. Scalars stay the degenerate single-slot case so existing behaviour is unchanged; a branch/body that changes a container's shape or an element's dtype now fails with a clear error. Adds MLIR-level unit tests and device system tests for list/tuple/dict/nested carry across if/for/while/ifexp. Co-Authored-By: Claude Opus 4.8 (1M context) --- python/flydsl/compiler/ast_rewriter.py | 322 +++++++++++------ .../test_dynamic_controlflow_list_carry.py | 342 ++++++++++++++++++ ...t_dynamic_controlflow_list_carry_device.py | 145 ++++++++ 3 files changed, 698 insertions(+), 111 deletions(-) create mode 100644 tests/unit/test_dynamic_controlflow_list_carry.py create mode 100644 tests/unit/test_dynamic_controlflow_list_carry_device.py 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/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 diff --git a/tests/unit/test_dynamic_controlflow_list_carry_device.py b/tests/unit/test_dynamic_controlflow_list_carry_device.py new file mode 100644 index 000000000..d4ed77879 --- /dev/null +++ b/tests/unit/test_dynamic_controlflow_list_carry_device.py @@ -0,0 +1,145 @@ +#!/usr/bin/env python3 + +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2025 FlyDSL Project Contributors + +"""System (device) tests: carry a Python list through dynamic if/for/while in a +real @flyc.kernel, launch it, and check the results against the expected values. + +This is the end-to-end counterpart of 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 flydsl.compiler as flyc +import flydsl.expr as fx + +pytestmark = [pytest.mark.l2_device, pytest.mark.rocm_lower] + +try: + import torch +except ImportError: + torch = None + +if torch is None or 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 TestDynamicControlFlowListCarryDevice: + 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() From c7450d4b32cbb5fdea2260a5b0cafac25bf1331f Mon Sep 17 00:00:00 2001 From: xudoyuan Date: Mon, 20 Jul 2026 10:36:02 +0000 Subject: [PATCH 2/4] [tests] move dynamic control-flow list-carry e2e test to tests/system The device/launch test belongs with the other end-to-end tests under tests/system (renamed to the *_e2e.py convention); the MLIR-level unit test stays under tests/unit. Aligns the header with the existing e2e tests (direct torch import + skip). Co-Authored-By: Claude Opus 4.8 (1M context) --- .../test_dynamic_controlflow_list_carry_e2e.py} | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename tests/{unit/test_dynamic_controlflow_list_carry_device.py => system/test_dynamic_controlflow_list_carry_e2e.py} (100%) diff --git a/tests/unit/test_dynamic_controlflow_list_carry_device.py b/tests/system/test_dynamic_controlflow_list_carry_e2e.py similarity index 100% rename from tests/unit/test_dynamic_controlflow_list_carry_device.py rename to tests/system/test_dynamic_controlflow_list_carry_e2e.py From 50f5181860803859dbdb3c3d4873e3fc6ef7fdbb Mon Sep 17 00:00:00 2001 From: xudoyuan Date: Mon, 20 Jul 2026 10:36:35 +0000 Subject: [PATCH 3/4] [tests] align list-carry e2e test with tests/system conventions Rename the test class to *E2E, reword the module docstring, and use the direct `import torch` + skip style shared by the other tests/system e2e files. Co-Authored-By: Claude Opus 4.8 (1M context) --- ...test_dynamic_controlflow_list_carry_e2e.py | 20 ++++++++----------- 1 file changed, 8 insertions(+), 12 deletions(-) diff --git a/tests/system/test_dynamic_controlflow_list_carry_e2e.py b/tests/system/test_dynamic_controlflow_list_carry_e2e.py index d4ed77879..8c5e5206e 100644 --- a/tests/system/test_dynamic_controlflow_list_carry_e2e.py +++ b/tests/system/test_dynamic_controlflow_list_carry_e2e.py @@ -3,28 +3,24 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright (c) 2025 FlyDSL Project Contributors -"""System (device) tests: carry a Python list through dynamic if/for/while in a +"""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 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``. +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] -try: - import torch -except ImportError: - torch = None - -if torch is None or not torch.cuda.is_available(): +if not torch.cuda.is_available(): pytest.skip("CUDA/ROCm not available", allow_module_level=True) @@ -107,7 +103,7 @@ def _run_ifexp_list(Out: fx.Tensor, flag: fx.Int32, stream: fx.Stream = fx.Strea # ── Tests ─────────────────────────────────────────────────────────────────── -class TestDynamicControlFlowListCarryDevice: +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 From f28c99cc339591120407e53dfa5c1256bee10fc0 Mon Sep 17 00:00:00 2001 From: xudoyuan Date: Fri, 24 Jul 2026 04:52:08 +0000 Subject: [PATCH 4/4] [compiler] unify container carry with protocol.py; add dict support Replace the bespoke explode/assemble recursion in the dynamic control-flow dispatchers with the existing protocol.extract_to_ir_values / construct_from_ir_values, so there is a single source of truth for the DSL-value <-> flat ir.Value round trip. Removes the duplicate _explode_carried / _assemble_carried / _carried_slot_count helpers (and the now-dead _unwrap_mlir_values / _pack_dispatch_results); keeps only the three shared unpack/pack/verify helpers reused by if/for/while/ifexp. protocol.py: add dict handling to get_ir_types / extract_to_ir_values / construct_from_ir_values (parallel to the SimpleNamespace branch); make the construct list/tuple branch dispatch on the exemplar so a top-level container works while staying backward compatible with the type-sequence call; add a bare-scalar coercion fallback to extract so python scalars still promote. Net -96 lines in ast_rewriter.py. unit + struct + dispatch + device e2e green. Co-Authored-By: Claude Opus 4.8 (1M context) --- python/flydsl/compiler/ast_rewriter.py | 182 ++++++------------------- python/flydsl/compiler/protocol.py | 36 ++++- 2 files changed, 74 insertions(+), 144 deletions(-) diff --git a/python/flydsl/compiler/ast_rewriter.py b/python/flydsl/compiler/ast_rewriter.py index 1eff34554..87c62fb6c 100644 --- a/python/flydsl/compiler/ast_rewriter.py +++ b/python/flydsl/compiler/ast_rewriter.py @@ -16,9 +16,9 @@ from .._mlir.dialects import arith, scf from ..expr import const_expr from ..expr.meta import capture_user_location, file_location -from ..expr.typing import as_dsl_value, as_ir_value +from ..expr.typing import as_ir_value from ..utils import env, log -from .protocol import construct_from_ir_values +from .protocol import construct_from_ir_values, extract_to_ir_values def _set_lineno(node, n=1): @@ -60,123 +60,47 @@ def _unwrap_constexpr(node): # --------------------------------------------------------------------------- -# Explode / assemble Python containers for control-flow carried state. +# Unpack / pack 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. +# nested) whose elements are DSL values. Following a pytree-style flatten / +# unflatten, we unpack such a container into a flat ir.Value list before +# building the scf op, and pack it back (during body tracing and on the results) +# so user code and downstream see the original container shape. # -# 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. +# There is no separate treedef object: the *pre-region value* itself is the +# structural template. The 3 helpers below are the only shared pieces; they are +# reused by all four control-flow dispatchers (if / for / while / ifexp) and +# delegate the actual flatten/rebuild to the existing DSL<->ir.Value protocol +# (protocol.extract_to_ir_values / construct_from_ir_values). Per-construct logic +# (building the specific scf op, then/else vs before/after wiring) stays inline +# in each dispatcher. # --------------------------------------------------------------------------- -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. + """UNPACK. Flatten every carried ``value`` (scalar or nested container) into + one shared flat ir.Value list via ``protocol.extract_to_ir_values``. - 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). + Returns ``(state_raw, counts, exemplars, result_types)`` where ``counts[i]`` + is how many ir.Values ``names[i]`` contributed and ``exemplars[i]`` is its + pre-region value (the template used later to pack the flat values back). + Called by every dispatcher before building the scf op. """ 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))." - ) + try: + slots = extract_to_ir_values(value) + except TypeError as exc: + raise TypeError( + f"{ctx_label} variable '{name}' of type {type(value).__name__} cannot be " + f"carried through dynamic control flow ({exc}). Its elements must be " + "MLIR-backed DSL values (e.g. fx.Int32(0), fx.Float32(0.0))." + ) from exc counts.append(len(slots)) exemplars.append(value) state_raw.extend(slots) @@ -185,37 +109,39 @@ def _explode_states(names, values, ctx_label): 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`.""" + """PACK (inverse of :func:`_explode_states`). Split flat ``slots`` per carried + name (by ``counts``) and rebuild each value from its ``exemplar`` template via + ``protocol.construct_from_ir_values``. Used to hand the region body its + containers (from block args) and to rebuild the final scf results. + """ 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) + out.append(construct_from_ir_values(type(exemplar), exemplar, seg)) 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.""" + """UNPACK + VERIFY. Flatten a region body/branch's produced values and check + their structure matches the region entry: same per-name ir.Value count AND + same per-slot dtype (a pytree-style structural equality check). + Returns the flat list to ``scf.yield``. Used by every dispatcher. + """ 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." + f"region: entered with {c_in} MLIR value(s) but the branch/body produced " + f"{c_out}. A carried container's length/nesting must match 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." + f"{ctx_label}: carried element dtype mismatch at slot {slot}: expected " + f"{t_in}, got {t_out}. A carried element's dtype must match on entry and exit." ) return out_raw @@ -709,28 +635,6 @@ def _normalize_branch_result(branch_result, state_names, state_map, branch_label ) return values - @staticmethod - def _unwrap_mlir_values(values, state_names, branch_label): - raw_values = [] - for name, value in zip(state_names, values): - raw = as_ir_value(value) - if not isinstance(raw, ir.Value): - raise TypeError( - f"if/else variable '{name}' in {branch_label} is {type(raw).__name__}, " - "not an MLIR Value. Only MLIR Values can be yielded from dynamic if/else branches." - ) - raw_values.append(raw) - return raw_values - - @staticmethod - def _pack_dispatch_results(results, state_values): - if not results: - return None - wrapped = [as_dsl_value(v, exemplar) for v, exemplar in zip(results, state_values)] - if len(wrapped) == 1: - return wrapped[0] - return tuple(wrapped) - @staticmethod def _collect_result_dict(result_names, local_vars): return {name: local_vars[name] for name in result_names} diff --git a/python/flydsl/compiler/protocol.py b/python/flydsl/compiler/protocol.py index 09fe1e0f1..045b3244d 100644 --- a/python/flydsl/compiler/protocol.py +++ b/python/flydsl/compiler/protocol.py @@ -41,6 +41,8 @@ def get_ir_types(obj) -> List[ir.Type]: return obj.__get_ir_types__() if isinstance(obj, SimpleNamespace): return list(chain.from_iterable(get_ir_types(v) for v in vars(obj).values())) + if isinstance(obj, dict): + return list(chain.from_iterable(get_ir_types(v) for v in obj.values())) if isinstance(obj, (tuple, list)): return list(chain.from_iterable(get_ir_types(x) for x in obj)) # derive IR types from extract_to_ir_values if possible @@ -82,8 +84,17 @@ def extract_to_ir_values(obj) -> List[ir.Value]: return obj.__extract_to_ir_values__() if isinstance(obj, SimpleNamespace): return list(chain.from_iterable(extract_to_ir_values(v) for v in vars(obj).values())) + if isinstance(obj, dict): + return list(chain.from_iterable(extract_to_ir_values(v) for v in obj.values())) if isinstance(obj, (tuple, list)): return list(chain.from_iterable(extract_to_ir_values(x) for x in obj)) + # Bare scalar leaf (e.g. a python int/float/bool, or an object exposing ir_value()): + # coerce to a single ir.Value so scalars can be carried like any other leaf. + from ..expr.typing import as_ir_value + + raw = as_ir_value(obj) + if isinstance(raw, ir.Value): + return [raw] raise TypeError(f"Cannot extract IR values from {obj}") @@ -99,16 +110,31 @@ def construct_from_ir_values(dsl_type, args, values: List[ir.Value]) -> DslType: if cursor != len(values): raise ValueError(f"SimpleNamespace expected {cursor} ir.Values, got {len(values)}") return SimpleNamespace(**rebuilt) + if isinstance(args, dict): + rebuilt = {} + cursor = 0 + for key, value in args.items(): + n = len(get_ir_types(value)) + rebuilt[key] = construct_from_ir_values(type(value), value, values[cursor : cursor + n]) + cursor += n + if cursor != len(values): + raise ValueError(f"dict expected {cursor} ir.Values, got {len(values)}") + return rebuilt if hasattr(dsl_type, "__construct_from_ir_values__"): exemplar = args if not isinstance(args, type) else None return dsl_type.__construct_from_ir_values__(values, exemplar) - if isinstance(dsl_type, (tuple, list)): + if isinstance(args, (tuple, list)): + # Dispatch on args (the exemplar), like the SimpleNamespace/dict branches, so a + # top-level list/tuple exemplar works too. When dsl_type is itself a sequence of + # types (e.g. the @jit param types) use it; otherwise derive per-element types. + types_seq = dsl_type if isinstance(dsl_type, (tuple, list)) else [type(a) for a in args] elems = [] - for ty, arg in zip(dsl_type, args, strict=True): + cursor = 0 + for ty, arg in zip(types_seq, args, strict=True): count = len(get_ir_types(arg)) - elems.append(construct_from_ir_values(ty, arg, values[:count])) - values = values[count:] - return type(dsl_type)(elems) + elems.append(construct_from_ir_values(ty, arg, values[cursor : cursor + count])) + cursor += count + return type(args)(elems) raise TypeError(f"Cannot construct DSL value for {dsl_type}")