From 1da5b7da33af9e3300e5574c5d42ee017472f7dc Mon Sep 17 00:00:00 2001 From: gf712 Date: Sun, 9 Aug 2026 20:26:42 +0100 Subject: [PATCH 01/12] runtime: Fix r0 clobber in import name execution --- src/executable/bytecode/instructions/ImportName.cpp | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/executable/bytecode/instructions/ImportName.cpp b/src/executable/bytecode/instructions/ImportName.cpp index 9aba3dca..234e5c90 100644 --- a/src/executable/bytecode/instructions/ImportName.cpp +++ b/src/executable/bytecode/instructions/ImportName.cpp @@ -22,6 +22,8 @@ PyResult ImportName::execute(VirtualMachine &vm, Interpreter &interpreter const auto &from_list = vm.reg(m_from_list); const auto &level = vm.reg(m_level); + [[maybe_unused]] RAIIStoreNonCallInstructionData non_call_instruction_data; + auto *builtins = interpreter.execution_frame()->builtins(); auto import_str = PyString::create("__import__"); if (import_str.is_err()) return import_str; @@ -68,4 +70,4 @@ std::vector ImportName::serialize() const }; return result; -} \ No newline at end of file +} From edaa4afec98c14e3aa92b9076be6e5c4c04b151e Mon Sep 17 00:00:00 2001 From: gf712 Date: Sun, 9 Aug 2026 20:42:50 +0100 Subject: [PATCH 02/12] runtime: fix the parsing of 4th arg in __import__ --- src/runtime/modules/BuiltinsModule.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/runtime/modules/BuiltinsModule.cpp b/src/runtime/modules/BuiltinsModule.cpp index 6aa54c8d..5732f41d 100644 --- a/src/runtime/modules/BuiltinsModule.cpp +++ b/src/runtime/modules/BuiltinsModule.cpp @@ -513,7 +513,7 @@ PyResult import(const PyTuple *args, const PyDict *, Interpreter &) auto *fromlist = arg3.unwrap(); auto arg4 = [args]() -> PyResult { - if (args->size() > 1) { + if (args->size() > 4) { auto arg4 = args->operator[](4); if (arg4.is_err()) return arg4; auto *level = arg4.unwrap(); From c86e78dddfbd8c674658469540ef6bd26f2b1cf8 Mon Sep 17 00:00:00 2001 From: gf712 Date: Mon, 10 Aug 2026 08:59:03 +0100 Subject: [PATCH 03/12] vm: fix instructions that clobber r0 --- .../bytecode/instructions/BuildDict.cpp | 29 ++++++++++++------- .../bytecode/instructions/BuildSet.cpp | 27 ++++++++++------- .../bytecode/instructions/FormatValue.cpp | 7 +++-- .../bytecode/instructions/ImportName.cpp | 7 +++-- .../bytecode/instructions/ListExtend.cpp | 11 +++++-- .../bytecode/instructions/SetUpdate.cpp | 4 +++ 6 files changed, 55 insertions(+), 30 deletions(-) diff --git a/src/executable/bytecode/instructions/BuildDict.cpp b/src/executable/bytecode/instructions/BuildDict.cpp index b8da29c2..eb06c9d7 100644 --- a/src/executable/bytecode/instructions/BuildDict.cpp +++ b/src/executable/bytecode/instructions/BuildDict.cpp @@ -6,19 +6,26 @@ using namespace py; PyResult BuildDict::execute(VirtualMachine &vm, Interpreter &) const { - PyDict::MapType map; - - if (m_size > 0) { - auto *start = vm.sp() - (m_size * 2); - for (size_t i = 0; i < m_size; ++i) { - const auto &key = *start; - const auto &value = *(start + m_size); - map.emplace(key, value); - start = std::next(start); + // Hashing the keys may run a Python __hash__/__eq__ and clobber r0. + auto dict_ = [&] { + [[maybe_unused]] RAIIStoreNonCallInstructionData non_call_instruction_data; + + PyDict::MapType map; + + if (m_size > 0) { + auto *start = vm.sp() - (m_size * 2); + for (size_t i = 0; i < m_size; ++i) { + const auto &key = *start; + const auto &value = *(start + m_size); + map.emplace(key, value); + start = std::next(start); + } } - } - return PyDict::create(map).and_then([&vm, this](PyDict *dict) { + return PyDict::create(map); + }(); + + return dict_.and_then([&vm, this](PyDict *dict) { vm.reg(m_dst) = dict; return Ok(dict); }); diff --git a/src/executable/bytecode/instructions/BuildSet.cpp b/src/executable/bytecode/instructions/BuildSet.cpp index 65f9ebdb..af35552b 100644 --- a/src/executable/bytecode/instructions/BuildSet.cpp +++ b/src/executable/bytecode/instructions/BuildSet.cpp @@ -6,16 +6,23 @@ using namespace py; PyResult BuildSet::execute(VirtualMachine &vm, Interpreter &) const { - PySet::SetType elements; - elements.reserve(m_size); - if (m_size > 0) { - auto *start = vm.sp() - m_size; - while (start != vm.sp()) { - elements.insert(*start); - start = std::next(start); + // Hashing the elements may run a Python __hash__/__eq__ and clobber r0. + auto set_ = [&] { + [[maybe_unused]] RAIIStoreNonCallInstructionData non_call_instruction_data; + + PySet::SetType elements; + elements.reserve(m_size); + if (m_size > 0) { + auto *start = vm.sp() - m_size; + while (start != vm.sp()) { + elements.insert(*start); + start = std::next(start); + } } - } - return PySet::create(elements).and_then([&vm, this](PySet *set) { + return PySet::create(elements); + }(); + + return set_.and_then([&vm, this](PySet *set) { vm.reg(m_dst) = set; return Ok(set); }); @@ -30,4 +37,4 @@ std::vector BuildSet::serialize() const m_dst, static_cast(m_size), }; -} \ No newline at end of file +} diff --git a/src/executable/bytecode/instructions/FormatValue.cpp b/src/executable/bytecode/instructions/FormatValue.cpp index bad7c8f9..97eab4b8 100644 --- a/src/executable/bytecode/instructions/FormatValue.cpp +++ b/src/executable/bytecode/instructions/FormatValue.cpp @@ -10,6 +10,10 @@ PyResult FormatValue::execute(VirtualMachine &vm, Interpreter &) const return PyObject::from(src) .and_then([this](PyObject *obj) { + // Every branch below can run a Python __str__/__repr__ and clobber + // r0. + [[maybe_unused]] RAIIStoreNonCallInstructionData non_call_instruction_data; + if (m_conversion == 0) { return obj->str(); } const auto conversion = static_cast(m_conversion); @@ -19,12 +23,9 @@ PyResult FormatValue::execute(VirtualMachine &vm, Interpreter &) const // return obj->ascii(); } break; case PyString::ReplacementField::Conversion::REPR: { - [[maybe_unused]] RAIIStoreNonCallInstructionData non_call_instruction_data; - return obj->repr(); } break; case PyString::ReplacementField::Conversion::STR: { - [[maybe_unused]] RAIIStoreNonCallInstructionData non_call_instruction_data; return obj->str(); } break; } diff --git a/src/executable/bytecode/instructions/ImportName.cpp b/src/executable/bytecode/instructions/ImportName.cpp index 234e5c90..ef05dfd3 100644 --- a/src/executable/bytecode/instructions/ImportName.cpp +++ b/src/executable/bytecode/instructions/ImportName.cpp @@ -22,8 +22,6 @@ PyResult ImportName::execute(VirtualMachine &vm, Interpreter &interpreter const auto &from_list = vm.reg(m_from_list); const auto &level = vm.reg(m_level); - [[maybe_unused]] RAIIStoreNonCallInstructionData non_call_instruction_data; - auto *builtins = interpreter.execution_frame()->builtins(); auto import_str = PyString::create("__import__"); if (import_str.is_err()) return import_str; @@ -51,7 +49,10 @@ PyResult ImportName::execute(VirtualMachine &vm, Interpreter &interpreter PyObject::from(level).unwrap()); if (args.is_err()) return args; - auto module = std::get(import_func)->call(args.unwrap(), nullptr); + auto module = [&] { + [[maybe_unused]] RAIIStoreNonCallInstructionData non_call_instruction_data; + return std::get(import_func)->call(args.unwrap(), nullptr); + }(); return module.and_then([&](auto *m) { vm.reg(m_destination) = m; diff --git a/src/executable/bytecode/instructions/ListExtend.cpp b/src/executable/bytecode/instructions/ListExtend.cpp index 4eb4b089..d871239c 100644 --- a/src/executable/bytecode/instructions/ListExtend.cpp +++ b/src/executable/bytecode/instructions/ListExtend.cpp @@ -16,8 +16,13 @@ PyResult ListExtend::execute(VirtualMachine &vm, Interpreter &) const ASSERT(pylist); ASSERT(as(pylist)); - return PyObject::from(value).and_then( - [pylist](PyObject *iterable) { return as(pylist)->extend(iterable); }); + return PyObject::from(value).and_then([pylist](PyObject *iterable) { + // extend() walks the iterator protocol, which may run Python __iter__ / + // __next__ and clobber r0. + [[maybe_unused]] RAIIStoreNonCallInstructionData non_call_instruction_data; + + return as(pylist)->extend(iterable); + }); } std::vector ListExtend::serialize() const @@ -27,4 +32,4 @@ std::vector ListExtend::serialize() const m_list, m_value, }; -} \ No newline at end of file +} diff --git a/src/executable/bytecode/instructions/SetUpdate.cpp b/src/executable/bytecode/instructions/SetUpdate.cpp index dd8941b2..0314f51d 100644 --- a/src/executable/bytecode/instructions/SetUpdate.cpp +++ b/src/executable/bytecode/instructions/SetUpdate.cpp @@ -15,6 +15,10 @@ PyResult SetUpdate::execute(VirtualMachine &vm, Interpreter &) const auto *pyset = std::get(set); ASSERT(pyset); ASSERT(as(pyset)); + // update() walks the iterator protocol and hashes each element, either of + // which may run Python and clobber r0. + [[maybe_unused]] RAIIStoreNonCallInstructionData non_call_instruction_data; + auto iterable_obj = PyObject::from(iterable); if (iterable_obj.is_err()) { return Err(iterable_obj.unwrap_err()); } From f1c504647c609250aaf34fba901dd22c9f558eb4 Mon Sep 17 00:00:00 2001 From: gf712 Date: Mon, 10 Aug 2026 11:06:32 +0100 Subject: [PATCH 04/12] mlir: fix lowering of default None args in methods --- .../FunctionPatterns.cpp | 27 ++++++++++++++----- 1 file changed, 20 insertions(+), 7 deletions(-) diff --git a/src/executable/mlir/Conversion/PythonToPythonBytecode/FunctionPatterns.cpp b/src/executable/mlir/Conversion/PythonToPythonBytecode/FunctionPatterns.cpp index 30bb80bf..a463dd21 100644 --- a/src/executable/mlir/Conversion/PythonToPythonBytecode/FunctionPatterns.cpp +++ b/src/executable/mlir/Conversion/PythonToPythonBytecode/FunctionPatterns.cpp @@ -176,15 +176,28 @@ namespace { }) != cell_names.end()) { - ASSERT(!class_returns.empty()); + ASSERT(class_returns.size() == 1); auto return_op = class_returns.front(); ASSERT(return_op->getParentOp() == op.getOperation()); - ASSERT(return_op.getValue().getDefiningOp()); - rewriter.setInsertionPoint(return_op.getValue().getDefiningOp()); - rewriter.replaceOpWithNewOp( - return_op.getValue().getDefiningOp(), - mlir::py::PyObjectType::get(getContext()), - mlir::StringRef{ "__class__" }); + + auto *defining_op = return_op.getValue().getDefiningOp(); + ASSERT(defining_op); + + if (mlir::isa(defining_op)) { + rewriter.setInsertionPoint(return_op); + auto class_cell = + rewriter.create(return_op.getLoc(), + mlir::py::PyObjectType::get(getContext()), + mlir::StringRef{ "__class__" }); + rewriter.modifyOpInPlace( + return_op, [&] { return_op.getValueMutable().assign(class_cell); }); + } else { + rewriter.setInsertionPoint(defining_op); + rewriter.replaceOpWithNewOp( + defining_op, + mlir::py::PyObjectType::get(getContext()), + mlir::StringRef{ "__class__" }); + } } } From cf0c19d559890b4b752e4bc02d1af62a15b9ca91 Mon Sep 17 00:00:00 2001 From: gf712 Date: Mon, 10 Aug 2026 16:46:47 +0100 Subject: [PATCH 05/12] mlir: add PyLoopOpInterface and centralise the br_yield parent set --- .../PythonToPythonBytecode.cpp | 138 +++++++++--------- .../mlir/Dialect/Python/IR/CMakeLists.txt | 4 + src/executable/mlir/Dialect/Python/IR/Ops.cpp | 10 ++ .../Dialect/Python/IR/PythonInterfaces.td | 61 ++++++++ .../mlir/Dialect/Python/IR/PythonOps.hpp | 2 + .../mlir/Dialect/Python/IR/PythonOps.td | 7 +- 6 files changed, 148 insertions(+), 74 deletions(-) create mode 100644 src/executable/mlir/Dialect/Python/IR/PythonInterfaces.td diff --git a/src/executable/mlir/Conversion/PythonToPythonBytecode/PythonToPythonBytecode.cpp b/src/executable/mlir/Conversion/PythonToPythonBytecode/PythonToPythonBytecode.cpp index c7eaab95..3e309962 100644 --- a/src/executable/mlir/Conversion/PythonToPythonBytecode/PythonToPythonBytecode.cpp +++ b/src/executable/mlir/Conversion/PythonToPythonBytecode/PythonToPythonBytecode.cpp @@ -52,6 +52,21 @@ namespace py { // PythonToPythonBytecodePass, and the dedicated single-pattern // passes that wrap the structural ones. + // The ops whose regions py.br_yield terminates, and which a structural + // pattern flattens into the enclosing CFG. A walk looking for the yields + // that belong to *one* such op must stop at any other one, because the + // nested op's own pattern owns everything inside it. + // + // This is BranchYieldOp's ParentOneOf list (PythonOps.td) — the sixth + // region-bearing python op, py.class, is excluded because its region + // becomes a separate function rather than being flattened in place. + bool is_flattened_region_op(mlir::Operation *op) + { + static_assert(mlir::py::BranchYieldOp::hasTrait< + mlir::OpTrait::HasParent::Impl>()); + return mlir::isa(op); + } + // Shared walker used by both ForLoopOpLowering and WhileOpLowering // to lower py.br_yield ops nested inside a loop body to cf.br ops // that target the right block (continue→condition / step, break→ @@ -73,20 +88,14 @@ namespace py { std::function callback = [&rewriter, continue_target, break_target, skip_op, &callback]( mlir::Operation *operation) { - if (auto loop = mlir::dyn_cast(operation)) { - if (loop.getOrelse().empty()) { return WalkResult::skip(); } - loop.getOrelse().walk(callback); - return WalkResult::skip(); - } - if (auto loop = mlir::dyn_cast(operation)) { - if (loop.getOrelse().empty()) { return WalkResult::skip(); } - loop.getOrelse().walk(callback); + if (auto loop = mlir::dyn_cast(operation)) { + auto &orelse = loop.getLoopOrelseRegion(); + if (orelse.empty()) { return WalkResult::skip(); } + orelse.walk(callback); return WalkResult::skip(); } auto yield_op = mlir::dyn_cast(operation); if (!yield_op) { return WalkResult::advance(); } - static_assert(mlir::py::BranchYieldOp::hasTrait::Impl>()); // Kindless yields under try/with/try-handler don't // participate in the loop's continue/break flow. if (!yield_op.getKind().has_value() @@ -351,15 +360,7 @@ namespace py { { if (region.empty()) { return; } region.walk([callback](mlir::Operation *childOp) { - static_assert(mlir::py::BranchYieldOp::hasTrait::Impl>()); - if (mlir::isa(childOp) - || mlir::isa(childOp) - || mlir::isa(childOp) - || mlir::isa(childOp) - || mlir::isa(childOp)) { - return WalkResult::skip(); - } + if (is_flattened_region_op(childOp)) { return WalkResult::skip(); } if (mlir::isa(childOp)) { // Both normal-completion (kindless) and loop-control // (break/continue) yields are surfaced; the callback @@ -712,59 +713,52 @@ namespace py { } }; - op.getBody().walk([&rewriter, - exit_block, - cleanup_block, - endBlock, - &emit_normal_exit]( - mlir::Operation *childOp) { - static_assert(mlir::py::BranchYieldOp::hasTrait::Impl>()); - if (mlir::isa(childOp) - || mlir::isa(childOp) - || mlir::isa(childOp) - || mlir::isa(childOp) - || mlir::isa(childOp)) { - return WalkResult::skip(); - } - if (auto op = mlir::dyn_cast(childOp)) { - rewriter.setInsertionPoint(op); - if (op.getCause()) { - rewriter.replaceOpWithNewOp( - op, op.getException(), op.getCause(), BlockRange{ cleanup_block }); - } else if (op.getException()) { - rewriter.replaceOpWithNewOp( - op, op.getException(), nullptr, BlockRange{ cleanup_block }); - } else { - rewriter.replaceOpWithNewOp( - op, BlockRange{ cleanup_block }); + op.getBody().walk( + [&rewriter, exit_block, cleanup_block, endBlock, &emit_normal_exit]( + mlir::Operation *childOp) { + if (is_flattened_region_op(childOp)) { return WalkResult::skip(); } + if (auto op = mlir::dyn_cast(childOp)) { + rewriter.setInsertionPoint(op); + if (op.getCause()) { + rewriter.replaceOpWithNewOp(op, + op.getException(), + op.getCause(), + BlockRange{ cleanup_block }); + } else if (op.getException()) { + rewriter.replaceOpWithNewOp( + op, op.getException(), nullptr, BlockRange{ cleanup_block }); + } else { + rewriter.replaceOpWithNewOp( + op, BlockRange{ cleanup_block }); + } + } else if (auto y = mlir::dyn_cast(childOp); + y && !y.getKind().has_value()) { + auto *current = y->getBlock(); + auto *next = rewriter.splitBlock(current, y->getIterator()); + rewriter.setInsertionPointToEnd(current); + rewriter.create( + y->getLoc()); + rewriter.create(y->getLoc(), exit_block); + rewriter.eraseBlock(next); + } else if (auto y = mlir::dyn_cast(childOp); + y && y.getKind().has_value()) { + // break/continue out of the with body: leave the + // exception handler, run __exit__, then hand the marker + // to the enclosing loop on a dedicated exit path. + auto *current = y->getBlock(); + auto *next = rewriter.splitBlock(current, y->getIterator()); + auto *lc_block = rewriter.createBlock(endBlock); + rewriter.setInsertionPointToEnd(current); + rewriter.create( + y->getLoc()); + rewriter.create(y->getLoc(), lc_block); + rewriter.setInsertionPointToStart(lc_block); + emit_normal_exit(); + rewriter.create(y->getLoc(), y.getKindAttr()); + rewriter.eraseBlock(next); } - } else if (auto y = mlir::dyn_cast(childOp); - y && !y.getKind().has_value()) { - auto *current = y->getBlock(); - auto *next = rewriter.splitBlock(current, y->getIterator()); - rewriter.setInsertionPointToEnd(current); - mlir::emitpybytecode::LeaveExceptionHandle::create(rewriter, y->getLoc()); - mlir::cf::BranchOp::create(rewriter, y->getLoc(), exit_block); - rewriter.eraseBlock(next); - } else if (auto y = mlir::dyn_cast(childOp); - y && y.getKind().has_value()) { - // break/continue out of the with body: leave the - // exception handler, run __exit__, then hand the marker - // to the enclosing loop on a dedicated exit path. - auto *current = y->getBlock(); - auto *next = rewriter.splitBlock(current, y->getIterator()); - auto *lc_block = rewriter.createBlock(endBlock); - rewriter.setInsertionPointToEnd(current); - mlir::emitpybytecode::LeaveExceptionHandle::create(rewriter, y->getLoc()); - mlir::cf::BranchOp::create(rewriter, y->getLoc(), lc_block); - rewriter.setInsertionPointToStart(lc_block); - emit_normal_exit(); - mlir::py::BranchYieldOp::create(rewriter, y->getLoc(), y.getKindAttr()); - rewriter.eraseBlock(next); - } - return WalkResult::advance(); - }); + return WalkResult::advance(); + }); rewriter.inlineRegionBefore(op.getBody(), endBlock); @@ -1028,4 +1022,4 @@ namespace py { } }// namespace py -}// namespace mlir \ No newline at end of file +}// namespace mlir diff --git a/src/executable/mlir/Dialect/Python/IR/CMakeLists.txt b/src/executable/mlir/Dialect/Python/IR/CMakeLists.txt index 71c6c9b3..f048252a 100644 --- a/src/executable/mlir/Dialect/Python/IR/CMakeLists.txt +++ b/src/executable/mlir/Dialect/Python/IR/CMakeLists.txt @@ -2,6 +2,10 @@ set(LLVM_TARGET_DEFINITIONS PythonAttributes.td) mlir_tablegen(PythonOpsEnums.h.inc -gen-enum-decls -dialect=python EXTRA_INCLUDES ${MLIR_INCLUDE_DIRS} ${PROJECT_SOURCE_DIR}/src/executable/mlir/Dialect) mlir_tablegen(PythonOpsEnums.cpp.inc -gen-enum-defs -dialect=python EXTRA_INCLUDES ${MLIR_INCLUDE_DIRS} ${PROJECT_SOURCE_DIR}/src/executable/mlir/Dialect) +set(LLVM_TARGET_DEFINITIONS PythonInterfaces.td) +mlir_tablegen(PythonInterfaces.h.inc -gen-op-interface-decls -dialect=python EXTRA_INCLUDES ${MLIR_INCLUDE_DIRS} ${PROJECT_SOURCE_DIR}/src/executable/mlir/Dialect) +mlir_tablegen(PythonInterfaces.cpp.inc -gen-op-interface-defs -dialect=python EXTRA_INCLUDES ${MLIR_INCLUDE_DIRS} ${PROJECT_SOURCE_DIR}/src/executable/mlir/Dialect) + set(LLVM_TARGET_DEFINITIONS PythonOps.td) mlir_tablegen(Dialect.h.inc -gen-dialect-decls -dialect=python EXTRA_INCLUDES ${MLIR_INCLUDE_DIRS} ${PROJECT_SOURCE_DIR}/src/executable/mlir/Dialect) mlir_tablegen(Dialect.cpp.inc -gen-dialect-defs -dialect=python EXTRA_INCLUDES ${MLIR_INCLUDE_DIRS} ${PROJECT_SOURCE_DIR}/src/executable/mlir/Dialect) diff --git a/src/executable/mlir/Dialect/Python/IR/Ops.cpp b/src/executable/mlir/Dialect/Python/IR/Ops.cpp index 31a7e8c1..9e64763f 100644 --- a/src/executable/mlir/Dialect/Python/IR/Ops.cpp +++ b/src/executable/mlir/Dialect/Python/IR/Ops.cpp @@ -17,6 +17,8 @@ #include "Python/IR/Dialect.cpp.inc" +#include "Python/IR/PythonInterfaces.cpp.inc" + namespace mlir { namespace py { namespace { @@ -263,6 +265,14 @@ namespace py { } }// namespace + mlir::Region &WhileOp::getLoopBodyRegion() { return getBody(); } + + mlir::Region &WhileOp::getLoopOrelseRegion() { return getOrelse(); } + + mlir::Region &ForLoopOp::getLoopBodyRegion() { return getBody(); } + + mlir::Region &ForLoopOp::getLoopOrelseRegion() { return getOrelse(); } + // Based on CIR loop interface implementation void WhileOp::getSuccessorRegions(mlir::RegionBranchPoint point, llvm::SmallVectorImpl ®ions) diff --git a/src/executable/mlir/Dialect/Python/IR/PythonInterfaces.td b/src/executable/mlir/Dialect/Python/IR/PythonInterfaces.td new file mode 100644 index 00000000..049031c0 --- /dev/null +++ b/src/executable/mlir/Dialect/Python/IR/PythonInterfaces.td @@ -0,0 +1,61 @@ +#ifndef PYTHON_INTERFACES +#define PYTHON_INTERFACES + +include "mlir/IR/OpBase.td" + +// Shared by the ops that model a Python loop statement: py.for_loop and +// py.while. +// +// What the passes actually need from a loop is not "is it a loop" but its two +// control-flow-distinct regions, because Python treats them differently: +// +// * the *body* is what `break`/`continue` inside it bind to. Its terminating +// py.br_yield ops (kindless for fallthrough, kinded for break/continue) are +// consumed by this loop's own lowering. +// +// * the *orelse* is NOT part of the loop. It runs once, after a normal exit, +// so a `break`/`continue` written there binds to the *enclosing* loop and +// must be left for that loop's pattern. Only the orelse's kindless +// completion yield belongs to this loop, as a branch to its exit. +// +// Before this interface existed both distinctions were spelled out as +// `isa` / `isa` pairs at each site, which is how the +// exemption in replace_loop_branch_yields came to cover try/with but not a +// nested loop's orelse — the two are indistinguishable in that form. +def PyLoopOpInterface : OpInterface<"PyLoopOpInterface"> { + let description = [{ + A Python loop statement, whose body owns the break/continue written + inside it and whose orelse belongs to the enclosing loop. + }]; + let cppNamespace = "::mlir::py"; + + let methods = [ + InterfaceMethod< + /*desc=*/[{ + The loop body. `break`/`continue` yields directly in this region + bind to this loop. + }], + /*retTy=*/"::mlir::Region &", + /*methodName=*/"getLoopBodyRegion" + >, + InterfaceMethod< + /*desc=*/[{ + The `else` clause, empty when the loop has none. Loop-control + yields here bind to the *enclosing* loop, not this one. + }], + /*retTy=*/"::mlir::Region &", + /*methodName=*/"getLoopOrelseRegion" + >, + ]; + + let extraClassDeclaration = [{ + // True when `region` is this loop's orelse. Distinguishing by region and + // not just by parent op is what keeps a loop's own body yields out of + // scope when a walk is looking for a nested loop's orelse. + bool isLoopOrelse(::mlir::Region *region) { + return region == &getLoopOrelseRegion(); + } + }]; +} + +#endif // PYTHON_INTERFACES diff --git a/src/executable/mlir/Dialect/Python/IR/PythonOps.hpp b/src/executable/mlir/Dialect/Python/IR/PythonOps.hpp index c694349b..c8206037 100644 --- a/src/executable/mlir/Dialect/Python/IR/PythonOps.hpp +++ b/src/executable/mlir/Dialect/Python/IR/PythonOps.hpp @@ -15,5 +15,7 @@ #include "Python/IR/PythonOpsEnums.h.inc" +#include "Python/IR/PythonInterfaces.h.inc" + #define GET_OP_CLASSES #include "Python/IR/Ops.h.inc" diff --git a/src/executable/mlir/Dialect/Python/IR/PythonOps.td b/src/executable/mlir/Dialect/Python/IR/PythonOps.td index 88449fda..c47714ff 100644 --- a/src/executable/mlir/Dialect/Python/IR/PythonOps.td +++ b/src/executable/mlir/Dialect/Python/IR/PythonOps.td @@ -1,4 +1,5 @@ include "Python/IR/PythonAttributes.td" +include "Python/IR/PythonInterfaces.td" include "Python/IR/PythonTypes.td" include "mlir/Interfaces/CallInterfaces.td" include "mlir/Interfaces/ControlFlowInterfaces.td" @@ -443,7 +444,8 @@ def UnpackExpandOp : Python_Op<"unpack_ex"> { Python_PyObjectType:$rest); } -def ForLoopOp : Python_Op<"for_loop", [DeclareOpInterfaceMethods]> { +def ForLoopOp : Python_Op<"for_loop", [DeclareOpInterfaceMethods, + DeclareOpInterfaceMethods]> { let summary = "For loop representation"; let arguments = (ins Python_PyObjectType:$iterable); @@ -457,7 +459,8 @@ def ForLoopOp : Python_Op<"for_loop", [DeclareOpInterfaceMethods]> { +def WhileOp : Python_Op<"while", [DeclareOpInterfaceMethods, + DeclareOpInterfaceMethods]> { let summary = "While loop representation"; let regions = (region AnyRegion:$condition, AnyRegion:$body, AnyRegion:$orelse); From b3b775bcd20957cf1e33291fdf9f95590b08f69c Mon Sep 17 00:00:00 2001 From: gf712 Date: Mon, 10 Aug 2026 16:49:47 +0100 Subject: [PATCH 06/12] mlir: fix while/else lowering and bind loop-else break/continue correctly `while ... else` never lowered. WhileOpLowering inlined the orelse region without rewriting its normal-completion py.br_yield, so the yield survived with no py.while parent left to satisfy its HasParent trait --- integration/tests/loop_else_break_binding.py | 109 ++++++++++ integration/tests/while_else.py | 97 +++++++++ src/executable/mlir/Conversion/Passes.td | 5 + .../PythonToPythonBytecode.cpp | 191 +++++++++++------- .../PythonToPythonBytecode.hpp | 15 +- src/executable/mlir/compile.cpp | 7 +- 6 files changed, 341 insertions(+), 83 deletions(-) create mode 100644 integration/tests/loop_else_break_binding.py create mode 100644 integration/tests/while_else.py diff --git a/integration/tests/loop_else_break_binding.py b/integration/tests/loop_else_break_binding.py new file mode 100644 index 00000000..ec0f7413 --- /dev/null +++ b/integration/tests/loop_else_break_binding.py @@ -0,0 +1,109 @@ +"""`break`/`continue` inside a loop's `else` binds to the *enclosing* loop. + +A loop's else clause is not part of its body, so Python binds loop control written +there to whatever loop encloses the whole statement. The lowering used to get this +wrong in two different ways: + + * `ForLoopOpLowering` rewrote its orelse's trailing yield without checking the + yield's kind, so the inner loop swallowed a `break` meant for the outer one — + silently running every outer iteration. + + * the enclosing loop's walker did claim the yield when the nested loop was a + `while`, but emitted the branch while that `py.while` was still unlowered, + producing a cross-region block reference the verifier rejects. + +Both are now handled by deferring: a loop refuses to lower while a nested loop +still holds a break/continue in its orelse, so the nested loop is flattened into +the enclosing region first and the branch is same-region by construction. That +handshake is also why both loop patterns share one pass. +""" + +# break in a nested for's else breaks the OUTER for. +log = [] +for outer in [1, 2, 3]: + log.append(outer) + for inner in []: + pass + else: + break +assert log == [1], log + +# Same with a while as the inner loop. +log = [] +for outer in [1, 2, 3]: + log.append(outer) + while False: + pass + else: + break +assert log == [1], log + +# continue in a nested loop's else continues the OUTER loop, skipping the rest +# of the outer body. +log = [] +for outer in [1, 2, 3]: + log.append(outer) + while False: + pass + else: + continue + log.append("after-must-not-run") +assert log == [1, 2, 3], log + +log = [] +for outer in [1, 2, 3]: + log.append(outer) + for inner in []: + pass + else: + continue + log.append("after-must-not-run") +assert log == [1, 2, 3], log + +# A while as the enclosing loop. +log = [] +n = 0 +while n < 3: + n += 1 + log.append(n) + for inner in []: + pass + else: + continue + log.append("after-must-not-run") +assert log == [1, 2, 3], log + +# The inner loop's own body break still binds to the inner loop, and the inner +# else is then skipped. +log = [] +for outer in [1, 2]: + for inner in [10, 20]: + log.append((outer, inner)) + break + else: + log.append("inner-else-must-not-run") + log.append(("after", outer)) +assert log == [(1, 10), ("after", 1), (2, 10), ("after", 2)], log + +# Three levels: the break binds to the loop enclosing the loop whose else it is, +# i.e. the middle one, so the outermost keeps iterating. +log = [] +for a in [1, 2]: + for b in [10, 20]: + log.append((a, b)) + for c in []: + pass + else: + break + log.append(("outer", a)) +assert log == [(1, 10), ("outer", 1), (2, 10), ("outer", 2)], log + +# An else that neither breaks nor continues still falls through to the exit. +log = [] +for outer in [1, 2]: + for inner in []: + pass + else: + log.append(("else", outer)) + log.append(("after", outer)) +assert log == [("else", 1), ("after", 1), ("else", 2), ("after", 2)], log diff --git a/integration/tests/while_else.py b/integration/tests/while_else.py new file mode 100644 index 00000000..ac9ea597 --- /dev/null +++ b/integration/tests/while_else.py @@ -0,0 +1,97 @@ +"""`while ... else` lowering. + +A while loop's orelse region ends in a `py.br_yield` marking normal completion. +`WhileOpLowering` used to inline that region without rewriting the yield, so it +survived with no `py.while` parent left to satisfy its HasParent trait, and every +`while/else` aborted during lowering. Nested inside a `for`, the enclosing loop's +walker claimed the yield instead and branched to the *for*'s continue target from +inside the still-unlowered while region — the "reference to block defined in +another region" verifier error that blocked `import re`. + +The nested case below is the shape from CPython's `sre_compile.py:471` +(`_generate_overlap_table`): a while/else inside a for, with the `break` in the +while *body*, which is what made `import re` fail. +""" + +out = [] + +# Normal exit runs the else. +i = 0 +while i < 3: + i += 1 +else: + out.append(("normal", i)) +assert out[-1] == ("normal", 3), out + +# A break in the body skips the else. +i = 0 +while i < 3: + i += 1 + break +else: + out.append("break-must-not-run") +assert out[-1] == ("normal", 3), out + +# A false condition still runs the else. +while False: + out.append("body-must-not-run") +else: + out.append("false-from-start") +assert out[-1] == "false-from-start", out + +# continue reaches the else via normal exit. +i = 0 +while i < 3: + i += 1 + continue +else: + out.append(("continue", i)) +assert out[-1] == ("continue", 3), out + + +def in_a_function(): + # Same shapes inside a function body, which lowers through a separate region. + seen = [] + n = 0 + while n < 2: + n += 1 + else: + seen.append(n) + while False: + pass + else: + seen.append("else") + return seen + + +assert in_a_function() == [2, "else"], in_a_function() + + +def generate_overlap_table(prefix): + # sre_compile._generate_overlap_table, the shape that blocked `import re`. + table = [0] * len(prefix) + for i in range(1, len(prefix)): + idx = table[i - 1] + while prefix[i] != prefix[idx]: + if idx == 0: + table[i] = 0 + break + idx = table[idx - 1] + else: + table[i] = idx + 1 + return table + + +assert generate_overlap_table("aab") == [0, 1, 0], generate_overlap_table("aab") +assert generate_overlap_table("abab") == [0, 0, 1, 2], generate_overlap_table("abab") +assert generate_overlap_table("aaaa") == [0, 1, 2, 3], generate_overlap_table("aaaa") +assert generate_overlap_table("abcd") == [0, 0, 0, 0], generate_overlap_table("abcd") + +# while/else nested in a for: the else runs on every iteration that exits normally. +nested = [] +for k in [1, 2, 3]: + while False: + pass + else: + nested.append(k) +assert nested == [1, 2, 3], nested diff --git a/src/executable/mlir/Conversion/Passes.td b/src/executable/mlir/Conversion/Passes.td index d652f058..368ebbf1 100644 --- a/src/executable/mlir/Conversion/Passes.td +++ b/src/executable/mlir/Conversion/Passes.td @@ -6,6 +6,11 @@ def ConvertPythonToPythonBytecode : Pass<"convert-python-to-pythonbytecode"> { let constructor = "mlir::py::createPythonToPythonBytecodePass()"; } +def ConvertPyLoops : Pass<"convert-py-loops"> { + let summary = "Lower py.for_loop and py.while together to emitpybytecode control flow"; + let constructor = "mlir::py::createConvertLoopsPass()"; +} + def ConvertPyForLoop : Pass<"convert-py-forloop"> { let summary = "Lower py.for_loop to emitpybytecode control flow"; let constructor = "mlir::py::createConvertForLoopPass()"; diff --git a/src/executable/mlir/Conversion/PythonToPythonBytecode/PythonToPythonBytecode.cpp b/src/executable/mlir/Conversion/PythonToPythonBytecode/PythonToPythonBytecode.cpp index 3e309962..1c792bf0 100644 --- a/src/executable/mlir/Conversion/PythonToPythonBytecode/PythonToPythonBytecode.cpp +++ b/src/executable/mlir/Conversion/PythonToPythonBytecode/PythonToPythonBytecode.cpp @@ -41,17 +41,6 @@ namespace mlir { namespace py { namespace { - // All non-region-bearing patterns now live in per-family files: - // {Arith, AttributeSubscript, Collection, ControlFlow, Function, - // Import, LoadStore}Patterns.cpp, registered via the - // populate*Patterns() entry points below. - // The shared DirectReplaceLowering / DirectReplaceRegisterName - // helpers and add_identifier* utilities moved to LoweringHelpers.hpp. - // What remains in this file are the four region-bearing structural - // patterns (ForLoop / While / Try / With), the - // PythonToPythonBytecodePass, and the dedicated single-pattern - // passes that wrap the structural ones. - // The ops whose regions py.br_yield terminates, and which a structural // pattern flattens into the enclosing CFG. A walk looking for the yields // that belong to *one* such op must stop at any other one, because the @@ -67,42 +56,61 @@ namespace py { return mlir::isa(op); } - // Shared walker used by both ForLoopOpLowering and WhileOpLowering - // to lower py.br_yield ops nested inside a loop body to cf.br ops - // that target the right block (continue→condition / step, break→ - // end). Nested loops are walked into their orelse regions only; - // the loop body itself is skipped because the nested loop will be - // lowered by its own pattern. + // True when `yield_op` is a loop-control (break/continue) yield that binds to + // the loop *enclosing* `loop` rather than to `loop` itself — i.e. it sits in + // `loop`'s orelse, which is not part of the loop body. + bool binds_to_enclosing_loop(mlir::py::PyLoopOpInterface loop, + mlir::py::BranchYieldOp yield_op) + { + return yield_op.getKind().has_value() && loop.isLoopOrelse(yield_op->getParentRegion()); + } + + // True when some loop nested in `region` still holds a break/continue that + // binds to the loop being lowered — i.e. one sitting in that nested loop's + // orelse. Such a yield cannot be rewritten yet: it lives in a region that has + // not been flattened, so branching it to our target block would be a + // cross-region block reference, which is invalid IR. + // + // The caller defers (fails the match) until the nested loop lowers and inlines + // the yield into our region, the same innermost-first trick TryOpLowering uses + // for nested trys. Terminates because the innermost such loop has nothing + // nested to wait on. + bool has_pending_nested_orelse_control(mlir::Region ®ion) + { + if (region.empty()) { return false; } + bool pending = false; + region.walk([&pending](mlir::Operation *op) { + auto loop = mlir::dyn_cast(op); + if (!loop) { return WalkResult::advance(); } + loop.getLoopOrelseRegion().walk( + [&pending, loop](mlir::py::BranchYieldOp yield_op) { + if (binds_to_enclosing_loop(loop, yield_op)) { pending = true; } + }); + // Only this loop's own orelse matters here; anything deeper is the + // nested loop's problem and it defers on it in turn. + return WalkResult::skip(); + }); + return pending; + } + + // Shared walker used by both ForLoopOpLowering and WhileOpLowering to lower + // the py.br_yield ops in a loop body to cf.br ops targeting the right block + // (continue→condition / step, break→end). // - // `skip_op` allows a caller to short-circuit on yield ops whose - // enclosing loop matches a specific predicate — ForLoopOpLowering - // uses this to ignore yields that bind to the *outer* for-loop's - // orelse (which shouldn't be lowered as part of the inner loop - // pass). + // It stops at every nested flattened-region op: those own their own yields, + // and a break/continue in a nested loop's orelse only becomes ours once that + // loop has been flattened into our region (see + // has_pending_nested_orelse_control, which is what makes that ordering hold). void replace_loop_branch_yields(mlir::PatternRewriter &rewriter, mlir::Region ®ion, mlir::Block *continue_target, - mlir::Block *break_target, - llvm::function_ref skip_op) + mlir::Block *break_target) { - std::function callback = - [&rewriter, continue_target, break_target, skip_op, &callback]( - mlir::Operation *operation) { - if (auto loop = mlir::dyn_cast(operation)) { - auto &orelse = loop.getLoopOrelseRegion(); - if (orelse.empty()) { return WalkResult::skip(); } - orelse.walk(callback); - return WalkResult::skip(); - } + region.walk( + [&rewriter, continue_target, break_target](mlir::Operation *operation) { + if (is_flattened_region_op(operation)) { return WalkResult::skip(); } auto yield_op = mlir::dyn_cast(operation); if (!yield_op) { return WalkResult::advance(); } - // Kindless yields under try/with/try-handler don't - // participate in the loop's continue/break flow. - if (!yield_op.getKind().has_value() - && mlir::isa(yield_op->getParentOp())) { - return WalkResult::advance(); - } - if (skip_op && skip_op(yield_op)) { return WalkResult::advance(); } rewriter.setInsertionPoint(yield_op); if (!yield_op.getKind().has_value() || yield_op.getKind().value() == py::LoopOpKind::continue_) { @@ -111,8 +119,29 @@ namespace py { rewriter.replaceOpWithNewOp(yield_op, break_target); } return WalkResult::advance(); - }; - region.walk(callback); + }); + } + + // Rewrites a loop orelse region's normal-completion (kindless) py.br_yield ops + // into branches to the loop's exit block. + // + // Only the kindless ones. A `break`/`continue` written in an orelse binds to + // the loop *enclosing* this one, so those are left in place: once this region + // is inlined they sit directly in the enclosing loop's body, where its own + // replace_loop_branch_yields claims them. + void replace_orelse_completion_yields(mlir::PatternRewriter &rewriter, + mlir::Region ®ion, + mlir::Block *exit_target) + { + if (region.empty()) { return; } + region.walk([&rewriter, exit_target](mlir::Operation *operation) { + if (is_flattened_region_op(operation)) { return WalkResult::skip(); } + auto yield_op = mlir::dyn_cast(operation); + if (!yield_op || yield_op.getKind().has_value()) { return WalkResult::advance(); } + rewriter.setInsertionPoint(yield_op); + rewriter.replaceOpWithNewOp(yield_op, exit_target); + return WalkResult::advance(); + }); } // Collects the loop-control (break/continue) kinds that appear directly @@ -208,6 +237,11 @@ namespace py { mlir::LogicalResult matchAndRewrite(mlir::py::ForLoopOp op, mlir::PatternRewriter &rewriter) const final { + // Lower innermost-first: a break/continue in a nested loop's orelse + // binds to us, but only becomes rewritable once that loop has been + // flattened into our region. + if (has_pending_nested_orelse_control(op.getBody())) { return failure(); } + auto *initBlock = rewriter.getInsertionBlock(); auto initPos = rewriter.getInsertionPoint(); @@ -251,15 +285,7 @@ namespace py { rewriter.inlineRegionBefore( op.getStep(), *op->getParentRegion(), endBlock->getIterator()); - // Skip yields whose enclosing for-loop sits inside an - // outer for-loop's orelse — those belong to the outer - // pattern's rewrite, not this one. - auto skip_orelse_yields = [](mlir::py::BranchYieldOp y) { - auto forloop_op = y->getParentOfType(); - return forloop_op && &forloop_op.getOrelse() == y->getParentRegion(); - }; - replace_loop_branch_yields( - rewriter, op.getBody(), for_iter_block, endBlock, skip_orelse_yields); + replace_loop_branch_yields(rewriter, op.getBody(), for_iter_block, endBlock); ASSERT(!op.getBody().empty()); auto *body_exit_block = &op.getBody().back(); @@ -267,15 +293,7 @@ namespace py { rewriter.inlineRegionBefore( op.getBody(), *op->getParentRegion(), endBlock->getIterator()); - if (!op.getOrelse().empty()) { - auto *orelse_exit_block = &op.getOrelse().back(); - ASSERT(orelse_exit_block->getTerminator()); - if (mlir::isa(orelse_exit_block->getTerminator())) { - rewriter.setInsertionPointToEnd(orelse_exit_block); - rewriter.replaceOpWithNewOp( - orelse_exit_block->getTerminator(), endBlock); - } - } + replace_orelse_completion_yields(rewriter, op.getOrelse(), endBlock); rewriter.inlineRegionBefore( op.getOrelse(), *op->getParentRegion(), endBlock->getIterator()); @@ -293,6 +311,10 @@ namespace py { mlir::LogicalResult matchAndRewrite(mlir::py::WhileOp op, mlir::PatternRewriter &rewriter) const final { + // See ForLoopOpLowering: innermost-first, so a nested loop's orelse + // break/continue is in our region before we try to retarget it. + if (has_pending_nested_orelse_control(op.getBody())) { return failure(); } + auto *initBlock = rewriter.getInsertionBlock(); auto initPos = rewriter.getInsertionPoint(); @@ -326,23 +348,13 @@ namespace py { rewriter.eraseOp(condition_op); rewriter.inlineRegionBefore(condition, endBlock); - replace_loop_branch_yields(rewriter, - op.getBody(), - &condition_start, - endBlock, - /*skip_op=*/{}); + replace_loop_branch_yields(rewriter, op.getBody(), &condition_start, endBlock); rewriter.inlineRegionBefore(op.getBody(), endBlock); - // if (!op.getOrelse().empty()) { - // auto *orelse_exit_block = &op.getOrelse().back(); - // ASSERT(orelse_exit_block->getTerminator()); - // if (mlir::isa(orelse_exit_block->getTerminator())) { - // rewriter.setInsertionPointToEnd(orelse_exit_block); - // rewriter.replaceOpWithNewOp( - // orelse_exit_block->getTerminator(), endBlock); - // } - // } + // Without this the orelse's completion yield survives region inlining + // with no py.while parent left to satisfy its HasParent trait. + replace_orelse_completion_yields(rewriter, op.getOrelse(), endBlock); rewriter.inlineRegionBefore(op.getOrelse(), endBlock); rewriter.eraseOp(op); @@ -837,8 +849,8 @@ namespace py { // (Python source dialect + EmitPythonBytecode target dialect); the // patterns also create cf::BranchOp / func::FuncOp internally, but // those dialects are already loaded by the time the pipeline runs. - template - struct SinglePatternConversionPass : public PassWrapper> + template + struct PatternConversionPass : public PassWrapper> { void getDependentDialects(DialectRegistry ®istry) const override { @@ -850,7 +862,7 @@ namespace py { void runOnOperation() final { mlir::RewritePatternSet patterns(&this->getContext()); - patterns.template add(&this->getContext()); + patterns.template add(&this->getContext()); GreedyRewriteConfig config; config.setStrictness(GreedyRewriteStrictness::AnyOp); @@ -862,11 +874,34 @@ namespace py { } }; + template + using SinglePatternConversionPass = PatternConversionPass; + + inline constexpr char kConvertLoopsArg[] = "convert-py-loops"; inline constexpr char kConvertForLoopArg[] = "convert-py-forloop"; inline constexpr char kConvertWhileLoopArg[] = "convert-py-while"; inline constexpr char kConvertTryArg[] = "convert-py-try"; inline constexpr char kConvertWithArg[] = "convert-py-with"; + // Both loop patterns must share one greedy driver. A break/continue in a + // nested loop's orelse binds to the enclosing loop and can only be retargeted + // after the nested loop is flattened, which the patterns arrange by deferring + // (see has_pending_nested_orelse_control). Deferral can only resolve if the + // pattern being waited on is available in the same run: with `for` and `while` + // in separate passes, a `for` containing a `while ... else: continue` would + // wait for a pattern that does not run until the next pass, and never lower. + // + // The single-pattern passes below stay registered for python-mlir-opt, where + // running one lowering in isolation is what the Conversion lit tests want. + struct ConvertLoopsPass + : public PatternConversionPass + { + MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(ConvertLoopsPass) + }; + struct ConvertForLoopPass : public SinglePatternConversionPass(); } + std::unique_ptr createConvertLoopsPass() { return std::make_unique(); } + std::unique_ptr createConvertTryPass() { return std::make_unique(); } std::unique_ptr createConvertWithPass() { return std::make_unique(); } diff --git a/src/executable/mlir/Conversion/PythonToPythonBytecode/PythonToPythonBytecode.hpp b/src/executable/mlir/Conversion/PythonToPythonBytecode/PythonToPythonBytecode.hpp index 432c4de8..0be37c71 100644 --- a/src/executable/mlir/Conversion/PythonToPythonBytecode/PythonToPythonBytecode.hpp +++ b/src/executable/mlir/Conversion/PythonToPythonBytecode/PythonToPythonBytecode.hpp @@ -9,10 +9,17 @@ class Pass; namespace py { std::unique_ptr createPythonToPythonBytecodePass(); - // Dedicated passes for the four region-bearing control-flow ops. - // Each runs only its own lowering pattern and is meant to slot into - // the pipeline ahead of the monolithic conversion pass, so that - // canonicalize / CSE can be interleaved between them. Plan step 18. + // Dedicated passes for the region-bearing control-flow ops. Each runs only its + // own lowering pattern(s) and is meant to slot into the pipeline ahead of the + // monolithic conversion pass, so that canonicalize / CSE can be interleaved + // between them. Plan step 18. + // + // Both loops share one pass: their patterns defer to each other so that a + // nested loop is flattened before the enclosing one retargets the break / + // continue in its orelse, and that handshake needs both patterns in the same + // greedy run. createConvertForLoopPass/createConvertWhileLoopPass remain for + // python-mlir-opt, which drives a single lowering at a time. + std::unique_ptr createConvertLoopsPass(); std::unique_ptr createConvertForLoopPass(); std::unique_ptr createConvertWhileLoopPass(); std::unique_ptr createConvertTryPass(); diff --git a/src/executable/mlir/compile.cpp b/src/executable/mlir/compile.cpp index f501d4a5..7276b557 100644 --- a/src/executable/mlir/compile.cpp +++ b/src/executable/mlir/compile.cpp @@ -77,8 +77,11 @@ std::shared_ptr compile(std::shared_ptr node, pm.addPass(::mlir::py::createConvertWithPass()); pm.addPass(::mlir::createCanonicalizerPass()); pm.addPass(::mlir::createCSEPass()); - pm.addPass(::mlir::py::createConvertForLoopPass()); - pm.addPass(::mlir::py::createConvertWhileLoopPass()); + // One pass for both loop kinds: a `break`/`continue` in a nested loop's orelse + // binds to the enclosing loop, and the patterns defer to each other so the + // nested loop is flattened first. Split across two passes, a `for` holding a + // `while ... else: continue` would defer on a pattern that had not run yet. + pm.addPass(::mlir::py::createConvertLoopsPass()); pm.addPass(::mlir::createCanonicalizerPass()); pm.addPass(::mlir::createCSEPass()); pm.addPass(::mlir::py::createPythonToPythonBytecodePass()); From 04fecef35db6f56d84c7c4aed43801567056a6f7 Mon Sep 17 00:00:00 2001 From: gf712 Date: Mon, 10 Aug 2026 16:50:56 +0100 Subject: [PATCH 07/12] mlir: declare RegionBranchOpInterface on py.with --- integration/tests/with_control_flow.py | 98 +++++++++++++++++++ src/executable/mlir/Dialect/Python/IR/Ops.cpp | 25 ++++- .../mlir/Dialect/Python/IR/PythonOps.td | 2 +- 3 files changed, 123 insertions(+), 2 deletions(-) create mode 100644 integration/tests/with_control_flow.py diff --git a/integration/tests/with_control_flow.py b/integration/tests/with_control_flow.py new file mode 100644 index 00000000..92946211 --- /dev/null +++ b/integration/tests/with_control_flow.py @@ -0,0 +1,98 @@ +"""`with` interacting with loop and exception control flow. + +py.with was the only flattened-region python op not declaring +RegionBranchOpInterface, even though py.br_yield already modelled the body -> +parent edge for it. Declaring it lets MLIR's region DCE and canonicalization +reason about the body's reachability, which is exactly the machinery that decides +whether a `break`/`continue` crossing the with boundary survives. These are the +shapes that exercise those edges; `with.py` covers __enter__/__exit__ protocol +details instead. +""" + + +class CM: + def __init__(self, log, name): + self.log = log + self.name = name + + def __enter__(self): + self.log.append(("enter", self.name)) + return self + + def __exit__(self, *args): + self.log.append(("exit", self.name)) + return False + + +# break and continue leaving a with body must still run __exit__. +log = [] +for i in [1, 2, 3]: + with CM(log, i): + if i == 1: + continue + if i == 3: + break + log.append(("body", i)) +assert log == [ + ("enter", 1), + ("exit", 1), + ("enter", 2), + ("body", 2), + ("exit", 2), + ("enter", 3), + ("exit", 3), +], log + +# with in a while body, whose else still runs on normal exit. +out = [] +n = 0 +while n < 2: + n += 1 + with CM(out, n): + out.append(("body", n)) +else: + out.append("else") +assert out == [ + ("enter", 1), + ("body", 1), + ("exit", 1), + ("enter", 2), + ("body", 2), + ("exit", 2), + "else", +], out + +# Nested with inside a try inside a loop: both __exit__ calls run, innermost +# first, before the handler. +deep = [] +for i in [1]: + try: + with CM(deep, "outer"): + with CM(deep, "inner"): + raise ValueError("x") + except ValueError: + deep.append("caught") +assert deep == [ + ("enter", "outer"), + ("enter", "inner"), + ("exit", "inner"), + ("exit", "outer"), + "caught", +], deep + + +def with_in_a_function(): + seen = [] + for i in [1, 2]: + with CM(seen, i): + if i == 2: + return seen + return seen + + +assert with_in_a_function() == [ + ("enter", 1), + ("exit", 1), + ("enter", 2), + ("exit", 2), +], with_in_a_function() diff --git a/src/executable/mlir/Dialect/Python/IR/Ops.cpp b/src/executable/mlir/Dialect/Python/IR/Ops.cpp index 9e64763f..3c63dcdc 100644 --- a/src/executable/mlir/Dialect/Python/IR/Ops.cpp +++ b/src/executable/mlir/Dialect/Python/IR/Ops.cpp @@ -383,7 +383,30 @@ namespace py { } } - void BranchYieldOp::getSuccessorRegions(llvm::ArrayRef, + // py.with has a single region and no loop back-edge: enter the body, then exit + // to the parent. The matching terminator edge (body -> parent) was already + // modelled by BranchYieldOp's WithOp case below; this is the op side of the + // same contract, which py.with was the only flattened-region op not to declare. + void WithOp::getSuccessorRegions(mlir::RegionBranchPoint point, + llvm::SmallVectorImpl ®ions) + { + if (point.isParent()) { + regions.emplace_back(&getBody()); + } else if (predecessor_region(point) == &getBody()) { + // Same exit-to-parent convention as TryOp: emplace the containing + // region rather than RegionSuccessor::parent(). + regions.emplace_back(getOperation()->getParentRegion()); + } else { + llvm_unreachable("unexpected branch origin"); + } + } + + mlir::ValueRange WithOp::getSuccessorInputs(mlir::RegionSuccessor successor) + { + return region_or_block_arguments(getOperation(), successor); + } + + void BranchYieldOp::getSuccessorRegions(llvm::ArrayRef operands, llvm::SmallVectorImpl ®ions) { static_assert(BranchYieldOp::hasTrait< diff --git a/src/executable/mlir/Dialect/Python/IR/PythonOps.td b/src/executable/mlir/Dialect/Python/IR/PythonOps.td index c47714ff..340edf9d 100644 --- a/src/executable/mlir/Dialect/Python/IR/PythonOps.td +++ b/src/executable/mlir/Dialect/Python/IR/PythonOps.td @@ -487,7 +487,7 @@ def TryOp : Python_Op<"try", [DeclareOpInterfaceMethods>:$handlers); } -def WithOp : Python_Op<"with"> { +def WithOp : Python_Op<"with", [DeclareOpInterfaceMethods]> { let arguments = (ins Variadic:$items); let regions = (region AnyRegion:$body); From e165e2493e6c5c2459ad6121b0a107bb13a18213 Mon Sep 17 00:00:00 2001 From: gf712 Date: Mon, 10 Aug 2026 21:05:20 +0100 Subject: [PATCH 08/12] mlir: stop emitting a block once a statement has terminated the block --- integration/tests/unreachable_code.py | 98 +++++++++++++++++++ .../mlir/Dialect/Python/MLIRGenerator.cpp | 46 +++++---- .../mlir/Dialect/Python/MLIRGenerator.hpp | 8 ++ 3 files changed, 135 insertions(+), 17 deletions(-) create mode 100644 integration/tests/unreachable_code.py diff --git a/integration/tests/unreachable_code.py b/integration/tests/unreachable_code.py new file mode 100644 index 00000000..848a3f4f --- /dev/null +++ b/integration/tests/unreachable_code.py @@ -0,0 +1,98 @@ +"""Statements following a terminator in the same suite. + +`break`, `continue` and `return` leave the builder's insertion point in a block +that already ends in a terminator, so MLIRGenerator used to append whatever came +next in the suite *after* that terminator: + + 'python.br_yield' op must be the last operation in the parent block + +The verifier rejected it, but MLIR's region DCE reached it first and segfaulted +(deleteDeadness reading a null terminator), so the diagnostic never mattered. +MLIRGenerator::codegen_statements now stops at the first statement that terminates the +block, which is also what unreachable code means. + +Reduced from sre_parse._parse, which is why `import re` crashed during lowering. +Nothing here asserts on the unreachable statements themselves — they cannot run; +the point is that the module compiles and the reachable behaviour is right. +""" + + +def after_break(values): + seen = [] + for v in values: + seen.append(v) + if v == 2: + break + seen.append("unreachable") + raise ValueError("unreachable") + return seen + + +assert after_break([1, 2, 3]) == [1, 2], after_break([1, 2, 3]) + + +def after_continue(values): + seen = [] + for v in values: + if v == 2: + continue + seen.append("unreachable") + seen.append(v) + return seen + + +assert after_continue([1, 2, 3]) == [1, 3], after_continue([1, 2, 3]) + + +def after_return(a): + return a + 1 + b = a * 2 + raise ValueError("unreachable") + + +assert after_return(1) == 2, after_return(1) + + +def after_break_in_while(a): + n = 0 + while True: + n += 1 + if n >= a: + break + n = 999 + raise ValueError("unreachable") + return n + + +assert after_break_in_while(3) == 3, after_break_in_while(3) + + +def after_break_in_try(values): + seen = [] + for v in values: + try: + seen.append(v) + if v == 2: + break + raise ValueError("unreachable") + except ValueError: + seen.append("caught") + return seen + + +assert after_break_in_try([1, 2, 3]) == [1, 2], after_break_in_try([1, 2, 3]) + + +def after_raise(a): + if a: + raise ValueError("boom") + a = 999 + return a + + +try: + after_raise(True) + raise AssertionError("should have raised") +except ValueError as e: + assert str(e) == "boom", str(e) +assert after_raise(False) is False diff --git a/src/executable/mlir/Dialect/Python/MLIRGenerator.cpp b/src/executable/mlir/Dialect/Python/MLIRGenerator.cpp index b8d7579b..86521dc3 100644 --- a/src/executable/mlir/Dialect/Python/MLIRGenerator.cpp +++ b/src/executable/mlir/Dialect/Python/MLIRGenerator.cpp @@ -1169,7 +1169,7 @@ ast::Value *MLIRGenerator::visit(const ast::ClassDefinition *node) m_context.builder(), node->name(), m_context.filename(), node->source_location())), node->source_location()); - for (const auto &el : node->body()) { el->codegen(this); } + codegen_statements(node->body()); mlir::cf::BranchOp::create(m_context.builder(), loc(m_context.builder(), m_context.filename(), node->body().back()->source_location()), @@ -1492,7 +1492,7 @@ ast::Value *MLIRGenerator::visit(const ast::For *node) mlir::py::BranchYieldOp::create(m_context.builder(), m_context.builder().getUnknownLoc()); m_context.builder().setInsertionPointToStart(&body_start); - for (const auto &el : node->body()) { el->codegen(this); } + codegen_statements(node->body()); if (m_context.builder().getInsertionBlock()->empty() || !m_context.builder() .getInsertionBlock() @@ -1503,7 +1503,7 @@ ast::Value *MLIRGenerator::visit(const ast::For *node) if (!node->orelse().empty()) { m_context.builder().setInsertionPointToStart(orelse_block); - for (const auto &el : node->orelse()) { el->codegen(this); } + codegen_statements(node->orelse()); if (m_context.builder().getBlock()->empty() || !m_context.builder().getBlock()->back().hasTrait()) { mlir::py::BranchYieldOp::create(m_context.builder(), @@ -1589,7 +1589,7 @@ ast::Value *MLIRGenerator::visit(const ast::If *node) orelse_block); m_context.builder().setInsertionPointToStart(if_block); - for (const auto &el : node->body()) { el->codegen(this); } + codegen_statements(node->body()); if (m_context.builder().getBlock()->empty() || !m_context.builder().getBlock()->back().hasTrait()) { mlir::cf::BranchOp::create(m_context.builder(), @@ -1598,7 +1598,7 @@ ast::Value *MLIRGenerator::visit(const ast::If *node) } if (!node->orelse().empty()) { m_context.builder().setInsertionPointToStart(orelse_block); - for (const auto &el : node->orelse()) { el->codegen(this); } + codegen_statements(node->orelse()); if (m_context.builder().getBlock()->empty() || !m_context.builder().getBlock()->back().hasTrait()) { mlir::cf::BranchOp::create(m_context.builder(), @@ -1806,6 +1806,20 @@ ast::Value *MLIRGenerator::visit(const ast::ListComp *node) node->source_location()); } +bool MLIRGenerator::current_block_is_terminated() const +{ + auto *block = m_context.builder().getBlock(); + return block && !block->empty() && block->back().hasTrait(); +} + +void MLIRGenerator::codegen_statements(const std::vector &statements) +{ + for (const auto &statement : statements) { + statement->codegen(this); + if (current_block_is_terminated()) { break; } + } +} + ast::Value *MLIRGenerator::visit(const ast::Module *m) { m_context.module()->setLoc(loc(m_context.builder(), m->filename(), SourceLocation{ 0, 0 })); @@ -1833,7 +1847,7 @@ ast::Value *MLIRGenerator::visit(const ast::Module *m) auto *entry_block = module_fn.addEntryBlock(); auto *exit_block = module_fn.addBlock(); m_context.builder().setInsertionPointToEnd(entry_block); - for (const auto &node : m->body()) { node->codegen(this); } + codegen_statements(m->body()); // If a program does not end with a terminator instruction, jump to the exit_block if (m_context.builder().getBlock()->empty() @@ -2553,7 +2567,7 @@ MLIRGenerator::MLIRValue *MLIRGenerator::make_function(const std::string &functi captures = collect_function_captures(mangled_name); builder.setInsertionPointToStart(&f.front()); - for (const auto &el : body) { el->codegen(this); } + codegen_statements(body); if (builder.getBlock()->empty() || !builder.getBlock()->back().hasTrait()) { @@ -2644,9 +2658,7 @@ ast::Value *MLIRGenerator::visit(const ast::Try *node) loc(m_context.builder(), m_context.filename(), node->source_location()), final_block); m_context.builder().setInsertionPointToEnd(final_block); - if (!node->finalbody().empty()) { - for (auto el : node->finalbody()) { el->codegen(this); } - } + if (!node->finalbody().empty()) { codegen_statements(node->finalbody()); } if (!m_context.builder().getBlock()->empty() && m_context.builder().getBlock()->back().hasTrait()) { m_context.builder().createBlock(current->getParent()); @@ -2654,7 +2666,7 @@ ast::Value *MLIRGenerator::visit(const ast::Try *node) }); m_context.builder().setInsertionPointToStart(&try_op.getBody().emplaceBlock()); - for (const auto &el : node->body()) { el->codegen(this); } + codegen_statements(node->body()); if (m_context.builder().getBlock()->empty() || !m_context.builder().getBlock()->back().hasTrait()) { mlir::py::BranchYieldOp::create(m_context.builder(), @@ -2699,7 +2711,7 @@ ast::Value *MLIRGenerator::visit(const ast::Try *node) handler->source_location()); } ClearExceptionBeforeReturn clear_exception_before_return{ scope() }; - for (auto el : handler->body()) { el->codegen(this); } + codegen_statements(handler->body()); if (m_context.builder().getBlock()->empty() || !m_context.builder() .getBlock() @@ -2713,7 +2725,7 @@ ast::Value *MLIRGenerator::visit(const ast::Try *node) if (!node->orelse().empty()) { m_context.builder().setInsertionPointToStart(&try_op.getOrelse().front()); - for (auto el : node->orelse()) { el->codegen(this); } + codegen_statements(node->orelse()); if (m_context.builder().getBlock()->empty() || !m_context.builder().getBlock()->back().hasTrait()) { mlir::py::BranchYieldOp::create(m_context.builder(), @@ -2726,7 +2738,7 @@ ast::Value *MLIRGenerator::visit(const ast::Try *node) if (!node->finalbody().empty()) { m_context.builder().setInsertionPointToStart(&try_op.getFinally().front()); - for (auto el : node->finalbody()) { el->codegen(this); } + codegen_statements(node->finalbody()); if (m_context.builder().getBlock()->empty() || !m_context.builder().getBlock()->back().hasTrait()) { mlir::py::BranchYieldOp::create(m_context.builder(), @@ -2819,7 +2831,7 @@ ast::Value *MLIRGenerator::visit(const ast::While *node) mlir::py::ConditionOp::create(m_context.builder(), test.getLoc(), test); m_context.builder().setInsertionPointToStart(body_start_block); - for (const auto &el : node->body()) { el->codegen(this); } + codegen_statements(node->body()); if (m_context.builder().getBlock()->empty() || !m_context.builder().getBlock()->back().hasTrait()) { @@ -2829,7 +2841,7 @@ ast::Value *MLIRGenerator::visit(const ast::While *node) if (!node->orelse().empty()) { m_context.builder().setInsertionPointToStart(orelse_block); - for (const auto &el : node->orelse()) { el->codegen(this); } + codegen_statements(node->orelse()); if (m_context.builder().getBlock()->empty() || !m_context.builder().getBlock()->back().hasTrait()) { mlir::py::BranchYieldOp::create(m_context.builder(), @@ -2895,7 +2907,7 @@ ast::Value *MLIRGenerator::visit(const ast::With *node) auto &body_start = with.getBody().emplaceBlock(); m_context.builder().setInsertionPointToStart(&body_start); - for (const auto &el : node->body()) { el->codegen(this); } + codegen_statements(node->body()); if (m_context.builder().getBlock()->empty() || !m_context.builder().getBlock()->back().hasTrait()) { mlir::py::BranchYieldOp::create(m_context.builder(), diff --git a/src/executable/mlir/Dialect/Python/MLIRGenerator.hpp b/src/executable/mlir/Dialect/Python/MLIRGenerator.hpp index fd749b09..2523777d 100644 --- a/src/executable/mlir/Dialect/Python/MLIRGenerator.hpp +++ b/src/executable/mlir/Dialect/Python/MLIRGenerator.hpp @@ -121,6 +121,14 @@ class MLIRGenerator : ast::CodeGenerator template MLIRValue *new_value(Args &&...args); + // Emits a statement list — a body, orelse, finalbody or handler body — stopping + // once a statement has terminated the current block. + void codegen_statements(const std::vector &statements); + + // True when the block being built already ends in a terminator, i.e. nothing + // further may be appended to it. + bool current_block_is_terminated() const; + void store_name(std::string_view name, MLIRValue *value, const SourceLocation &location); MLIRValue *load_name(std::string_view name, const SourceLocation &location); void delete_name(std::string_view name, const SourceLocation &location); From a486da7060dd1adedaf630afb5bb248366d2c964 Mon Sep 17 00:00:00 2001 From: gf712 Date: Mon, 10 Aug 2026 21:05:36 +0100 Subject: [PATCH 09/12] mlir: build the while test where py.condition is, not at its operand's definition --- integration/tests/while_condition_cse.py | 93 +++++++++++++++++++ .../PythonToPythonBytecode.cpp | 10 +- 2 files changed, 96 insertions(+), 7 deletions(-) create mode 100644 integration/tests/while_condition_cse.py diff --git a/integration/tests/while_condition_cse.py b/integration/tests/while_condition_cse.py new file mode 100644 index 00000000..c4bc9b6e --- /dev/null +++ b/integration/tests/while_condition_cse.py @@ -0,0 +1,93 @@ +"""A while condition whose value is defined outside the condition region. + +WhileOpLowering built the loop's test and cf.cond_br at the *condition value's* +definition site. That is usually inside the condition region, but not always: CSE +merges the constant behind `while True:` with an identical constant in the +enclosing function, after which py.condition tests a value defined in the +function's entry block. Inserting there put the cf.cond_br in the middle of that +block, as a second terminator, and MLIR's region DCE then segfaulted on the block +whose last operation was no longer a terminator. + +py.condition is by construction the terminator of the condition region's last +block, and the value it tests necessarily dominates it, so that is where the +branch belongs. + +`b = True` before the loop is what creates the constant CSE merges with — without +it the loop's `True` is unique and the bug does not appear. Reduced from +sre_parse._parse; the same fault was the long-standing `import weakref` crash. +""" + + +def only_exit_is_raise(a): + b = True + if a: + while True: + raise ValueError("boom") + return b + + +try: + only_exit_is_raise(True) + raise AssertionError("should have raised") +except ValueError as e: + assert str(e) == "boom", str(e) +assert only_exit_is_raise(False) is True + + +def shared_true_constant(limit): + flag = True + n = 0 + while True: + n += 1 + if n >= limit: + break + return (n, flag) + + +assert shared_true_constant(3) == (3, True), shared_true_constant(3) + + +def shared_false_constant(a): + flag = False + n = 0 + while not flag: + n += 1 + if n >= a: + flag = True + return n + + +assert shared_false_constant(2) == 2, shared_false_constant(2) + + +def condition_is_a_parameter(cond, limit): + # The condition value is a block argument rather than an op result, the other + # branch of the insertion-point choice that used to exist. + n = 0 + while cond: + n += 1 + if n >= limit: + cond = False + return n + + +assert condition_is_a_parameter(True, 2) == 2, condition_is_a_parameter(True, 2) +assert condition_is_a_parameter(False, 2) == 0, condition_is_a_parameter(False, 2) + + +def nested_loops_sharing_true(limit): + t = True + outer = 0 + while True: + outer += 1 + inner = 0 + while True: + inner += 1 + if inner >= 2: + break + if outer >= limit: + break + return (outer, inner, t) + + +assert nested_loops_sharing_true(2) == (2, 2, True), nested_loops_sharing_true(2) diff --git a/src/executable/mlir/Conversion/PythonToPythonBytecode/PythonToPythonBytecode.cpp b/src/executable/mlir/Conversion/PythonToPythonBytecode/PythonToPythonBytecode.cpp index 1c792bf0..21efd74b 100644 --- a/src/executable/mlir/Conversion/PythonToPythonBytecode/PythonToPythonBytecode.cpp +++ b/src/executable/mlir/Conversion/PythonToPythonBytecode/PythonToPythonBytecode.cpp @@ -332,13 +332,9 @@ namespace py { rewriter.setInsertionPointToEnd(initBlock); mlir::cf::BranchOp::create(rewriter, condition_op.getLoc(), &condition_start); - if (mlir::isa(condition_op.getCond())) { - rewriter.setInsertionPointToStart(condition_op.getCond().getParentBlock()); - } else { - rewriter.setInsertionPointAfter(condition_op.getCond().getDefiningOp()); - } - auto should_jump = mlir::py::CastToBoolOp::create( - rewriter, condition_op.getLoc(), rewriter.getI1Type(), condition_op.getCond()); + rewriter.setInsertionPoint(condition_op); + auto should_jump = rewriter.create( + condition_op.getLoc(), rewriter.getI1Type(), condition_op.getCond()); ASSERT(!op.getBody().empty()); mlir::cf::CondBranchOp::create(rewriter, condition_op.getLoc(), From b85a94db8b07dd426c5b657e3a33c9ce4c28b785 Mon Sep 17 00:00:00 2001 From: gf712 Date: Sun, 23 Aug 2026 14:09:23 +0100 Subject: [PATCH 10/12] mlir: bind loop-else break/continue transitively --- integration/tests/loop_else_break_binding.py | 83 ++++++++++++++++++- .../PythonToPythonBytecode.cpp | 67 +++++++++------ 2 files changed, 121 insertions(+), 29 deletions(-) diff --git a/integration/tests/loop_else_break_binding.py b/integration/tests/loop_else_break_binding.py index ec0f7413..2ab8d4fa 100644 --- a/integration/tests/loop_else_break_binding.py +++ b/integration/tests/loop_else_break_binding.py @@ -13,9 +13,13 @@ producing a cross-region block reference the verifier rejects. Both are now handled by deferring: a loop refuses to lower while a nested loop -still holds a break/continue in its orelse, so the nested loop is flattened into -the enclosing region first and the branch is same-region by construction. That -handshake is also why both loop patterns share one pass. +still holds a break/continue that binds to it, so the nested loop is flattened +into the enclosing region first and the branch is same-region by construction. +That handshake is also why both loop patterns share one pass. + +Binding outwards is transitive: an else nested inside another else is still +lexically part of whichever loop body encloses the pair, so the deferral has to +follow the whole chain rather than stop one level in. """ # break in a nested for's else breaks the OUTER for. @@ -107,3 +111,76 @@ log.append(("else", outer)) log.append(("after", outer)) assert log == [("else", 1), ("after", 1), ("else", 2), ("after", 2)], log + +# An else nested inside another else: the break binds outwards through *both*, to +# the outermost loop. Checking only the first nested loop's else missed this and +# silently dropped the break. +log = [] +for a in [1, 2, 3]: + log.append(a) + for b in []: + pass + else: + for c in []: + pass + else: + break +assert log == [1], log + +# Same shape with `continue`, which must skip the rest of the outermost body. +log = [] +for a in [1, 2, 3]: + log.append(a) + for b in []: + pass + else: + for c in []: + pass + else: + continue + log.append("after-must-not-run") +assert log == [1, 2, 3], log + +# Same shape built from `while`s. Here the mis-binding was not silent: the branch +# was emitted while the inner py.while was still a region of its own, which the +# verifier rejects as a reference to a block in another region. +log = [] +for a in [1, 2, 3]: + log.append(a) + while False: + pass + else: + while False: + pass + else: + break +assert log == [1], log + +# Three elses deep, to check the walk follows the chain rather than a fixed depth. +log = [] +for a in [1, 2, 3]: + log.append(a) + for b in []: + pass + else: + for c in []: + pass + else: + for d in []: + pass + else: + break +assert log == [1], log + +# The chain stops at the first loop *body*: this break is in the body of a loop +# that happens to sit in an else, so it binds to that loop and no further. +log = [] +for a in [1, 2]: + for b in []: + pass + else: + for c in [10, 20]: + log.append((a, c)) + break + log.append(("after", a)) +assert log == [(1, 10), ("after", 1), (2, 10), ("after", 2)], log diff --git a/src/executable/mlir/Conversion/PythonToPythonBytecode/PythonToPythonBytecode.cpp b/src/executable/mlir/Conversion/PythonToPythonBytecode/PythonToPythonBytecode.cpp index 21efd74b..ece10c93 100644 --- a/src/executable/mlir/Conversion/PythonToPythonBytecode/PythonToPythonBytecode.cpp +++ b/src/executable/mlir/Conversion/PythonToPythonBytecode/PythonToPythonBytecode.cpp @@ -56,39 +56,54 @@ namespace py { return mlir::isa(op); } - // True when `yield_op` is a loop-control (break/continue) yield that binds to - // the loop *enclosing* `loop` rather than to `loop` itself — i.e. it sits in - // `loop`'s orelse, which is not part of the loop body. - bool binds_to_enclosing_loop(mlir::py::PyLoopOpInterface loop, - mlir::py::BranchYieldOp yield_op) + // True when `yield_op`, a loop-control (break/continue) yield, binds to the + // loop whose body region is `body`. + // + // Python binds break/continue to the innermost loop whose *body* lexically + // contains it. An else clause is not part of its own loop's body, so a yield + // sitting there keeps searching outwards — and transitively so: an else nested + // inside another else is still lexically part of whatever body encloses the + // pair. Regions that are neither body nor orelse (a try body, a with body) are + // likewise transparent, which is what makes `break` inside a `try` bind to the + // loop around it. + bool binds_to_loop(mlir::Region &body, mlir::py::BranchYieldOp yield_op) { - return yield_op.getKind().has_value() && loop.isLoopOrelse(yield_op->getParentRegion()); + for (mlir::Region *region = yield_op->getParentRegion(); region != nullptr; + region = region->getParentRegion()) { + if (region == &body) { return true; } + auto loop = + mlir::dyn_cast_if_present(region->getParentOp()); + // A loop body (or a for's step) stops the search: the yield is that + // loop's, and its own pattern claims it. + if (loop && !loop.isLoopOrelse(region)) { return false; } + } + return false; } - // True when some loop nested in `region` still holds a break/continue that - // binds to the loop being lowered — i.e. one sitting in that nested loop's - // orelse. Such a yield cannot be rewritten yet: it lives in a region that has - // not been flattened, so branching it to our target block would be a - // cross-region block reference, which is invalid IR. + // True when a break/continue that binds to the loop whose body is `body` is + // somewhere replace_loop_branch_yields cannot reach yet — inside a nested + // region that has not been flattened into ours. Branching it to our target + // block now would be a cross-region block reference, which is invalid IR. // - // The caller defers (fails the match) until the nested loop lowers and inlines + // The caller defers (fails the match) until the nested op lowers and inlines // the yield into our region, the same innermost-first trick TryOpLowering uses - // for nested trys. Terminates because the innermost such loop has nothing - // nested to wait on. - bool has_pending_nested_orelse_control(mlir::Region ®ion) + // for nested trys. Terminates because the innermost such op has nothing nested + // to wait on. + bool has_pending_nested_orelse_control(mlir::Region &body) { - if (region.empty()) { return false; } + if (body.empty()) { return false; } bool pending = false; - region.walk([&pending](mlir::Operation *op) { - auto loop = mlir::dyn_cast(op); - if (!loop) { return WalkResult::advance(); } - loop.getLoopOrelseRegion().walk( - [&pending, loop](mlir::py::BranchYieldOp yield_op) { - if (binds_to_enclosing_loop(loop, yield_op)) { pending = true; } - }); - // Only this loop's own orelse matters here; anything deeper is the - // nested loop's problem and it defers on it in turn. - return WalkResult::skip(); + body.walk([&pending, &body](mlir::Operation *op) { + // Mirror replace_loop_branch_yields: what it walks through it rewrites + // in place, so only what it skips over can be pending. + if (!is_flattened_region_op(op)) { return WalkResult::advance(); } + op->walk([&pending, &body](mlir::py::BranchYieldOp yield_op) { + if (!yield_op.getKind().has_value()) { return WalkResult::advance(); } + if (!binds_to_loop(body, yield_op)) { return WalkResult::advance(); } + pending = true; + return WalkResult::interrupt(); + }); + return pending ? WalkResult::interrupt() : WalkResult::skip(); }); return pending; } From 0037a3319e7e1ead774bc628667b1b3d45edb25d Mon Sep 17 00:00:00 2001 From: gf712 Date: Sun, 23 Aug 2026 14:21:45 +0100 Subject: [PATCH 11/12] build: fix deprecated mlir function --- .../FunctionPatterns.cpp | 8 +- .../PythonToPythonBytecode.cpp | 85 ++++++++----------- src/executable/mlir/Dialect/Python/IR/Ops.cpp | 2 +- 3 files changed, 40 insertions(+), 55 deletions(-) diff --git a/src/executable/mlir/Conversion/PythonToPythonBytecode/FunctionPatterns.cpp b/src/executable/mlir/Conversion/PythonToPythonBytecode/FunctionPatterns.cpp index a463dd21..440717ee 100644 --- a/src/executable/mlir/Conversion/PythonToPythonBytecode/FunctionPatterns.cpp +++ b/src/executable/mlir/Conversion/PythonToPythonBytecode/FunctionPatterns.cpp @@ -185,10 +185,10 @@ namespace { if (mlir::isa(defining_op)) { rewriter.setInsertionPoint(return_op); - auto class_cell = - rewriter.create(return_op.getLoc(), - mlir::py::PyObjectType::get(getContext()), - mlir::StringRef{ "__class__" }); + auto class_cell = mlir::emitpybytecode::LoadClosureOp::create(rewriter, + return_op.getLoc(), + mlir::py::PyObjectType::get(getContext()), + mlir::StringRef{ "__class__" }); rewriter.modifyOpInPlace( return_op, [&] { return_op.getValueMutable().assign(class_cell); }); } else { diff --git a/src/executable/mlir/Conversion/PythonToPythonBytecode/PythonToPythonBytecode.cpp b/src/executable/mlir/Conversion/PythonToPythonBytecode/PythonToPythonBytecode.cpp index ece10c93..3144bf6c 100644 --- a/src/executable/mlir/Conversion/PythonToPythonBytecode/PythonToPythonBytecode.cpp +++ b/src/executable/mlir/Conversion/PythonToPythonBytecode/PythonToPythonBytecode.cpp @@ -56,54 +56,39 @@ namespace py { return mlir::isa(op); } - // True when `yield_op`, a loop-control (break/continue) yield, binds to the - // loop whose body region is `body`. - // - // Python binds break/continue to the innermost loop whose *body* lexically - // contains it. An else clause is not part of its own loop's body, so a yield - // sitting there keeps searching outwards — and transitively so: an else nested - // inside another else is still lexically part of whatever body encloses the - // pair. Regions that are neither body nor orelse (a try body, a with body) are - // likewise transparent, which is what makes `break` inside a `try` bind to the - // loop around it. - bool binds_to_loop(mlir::Region &body, mlir::py::BranchYieldOp yield_op) + // True when `yield_op` is a loop-control (break/continue) yield that binds to + // the loop *enclosing* `loop` rather than to `loop` itself — i.e. it sits in + // `loop`'s orelse, which is not part of the loop body. + bool binds_to_enclosing_loop(mlir::py::PyLoopOpInterface loop, + mlir::py::BranchYieldOp yield_op) { - for (mlir::Region *region = yield_op->getParentRegion(); region != nullptr; - region = region->getParentRegion()) { - if (region == &body) { return true; } - auto loop = - mlir::dyn_cast_if_present(region->getParentOp()); - // A loop body (or a for's step) stops the search: the yield is that - // loop's, and its own pattern claims it. - if (loop && !loop.isLoopOrelse(region)) { return false; } - } - return false; + return yield_op.getKind().has_value() && loop.isLoopOrelse(yield_op->getParentRegion()); } - // True when a break/continue that binds to the loop whose body is `body` is - // somewhere replace_loop_branch_yields cannot reach yet — inside a nested - // region that has not been flattened into ours. Branching it to our target - // block now would be a cross-region block reference, which is invalid IR. + // True when some loop nested in `region` still holds a break/continue that + // binds to the loop being lowered — i.e. one sitting in that nested loop's + // orelse. Such a yield cannot be rewritten yet: it lives in a region that has + // not been flattened, so branching it to our target block would be a + // cross-region block reference, which is invalid IR. // - // The caller defers (fails the match) until the nested op lowers and inlines + // The caller defers (fails the match) until the nested loop lowers and inlines // the yield into our region, the same innermost-first trick TryOpLowering uses - // for nested trys. Terminates because the innermost such op has nothing nested - // to wait on. - bool has_pending_nested_orelse_control(mlir::Region &body) + // for nested trys. Terminates because the innermost such loop has nothing + // nested to wait on. + bool has_pending_nested_orelse_control(mlir::Region ®ion) { - if (body.empty()) { return false; } + if (region.empty()) { return false; } bool pending = false; - body.walk([&pending, &body](mlir::Operation *op) { - // Mirror replace_loop_branch_yields: what it walks through it rewrites - // in place, so only what it skips over can be pending. - if (!is_flattened_region_op(op)) { return WalkResult::advance(); } - op->walk([&pending, &body](mlir::py::BranchYieldOp yield_op) { - if (!yield_op.getKind().has_value()) { return WalkResult::advance(); } - if (!binds_to_loop(body, yield_op)) { return WalkResult::advance(); } - pending = true; - return WalkResult::interrupt(); - }); - return pending ? WalkResult::interrupt() : WalkResult::skip(); + region.walk([&pending](mlir::Operation *op) { + auto loop = mlir::dyn_cast(op); + if (!loop) { return WalkResult::advance(); } + loop.getLoopOrelseRegion().walk( + [&pending, loop](mlir::py::BranchYieldOp yield_op) { + if (binds_to_enclosing_loop(loop, yield_op)) { pending = true; } + }); + // Only this loop's own orelse matters here; anything deeper is the + // nested loop's problem and it defers on it in turn. + return WalkResult::skip(); }); return pending; } @@ -348,8 +333,8 @@ namespace py { mlir::cf::BranchOp::create(rewriter, condition_op.getLoc(), &condition_start); rewriter.setInsertionPoint(condition_op); - auto should_jump = rewriter.create( - condition_op.getLoc(), rewriter.getI1Type(), condition_op.getCond()); + auto should_jump = mlir::py::CastToBoolOp::create( + rewriter, condition_op.getLoc(), rewriter.getI1Type(), condition_op.getCond()); ASSERT(!op.getBody().empty()); mlir::cf::CondBranchOp::create(rewriter, condition_op.getLoc(), @@ -759,9 +744,9 @@ namespace py { auto *current = y->getBlock(); auto *next = rewriter.splitBlock(current, y->getIterator()); rewriter.setInsertionPointToEnd(current); - rewriter.create( - y->getLoc()); - rewriter.create(y->getLoc(), exit_block); + mlir::emitpybytecode::LeaveExceptionHandle::create( + rewriter, y->getLoc()); + mlir::cf::BranchOp::create(rewriter, y->getLoc(), exit_block); rewriter.eraseBlock(next); } else if (auto y = mlir::dyn_cast(childOp); y && y.getKind().has_value()) { @@ -772,12 +757,12 @@ namespace py { auto *next = rewriter.splitBlock(current, y->getIterator()); auto *lc_block = rewriter.createBlock(endBlock); rewriter.setInsertionPointToEnd(current); - rewriter.create( - y->getLoc()); - rewriter.create(y->getLoc(), lc_block); + mlir::emitpybytecode::LeaveExceptionHandle::create( + rewriter, y->getLoc()); + mlir::cf::BranchOp::create(rewriter, y->getLoc(), lc_block); rewriter.setInsertionPointToStart(lc_block); emit_normal_exit(); - rewriter.create(y->getLoc(), y.getKindAttr()); + mlir::py::BranchYieldOp::create(rewriter, y->getLoc(), y.getKindAttr()); rewriter.eraseBlock(next); } return WalkResult::advance(); diff --git a/src/executable/mlir/Dialect/Python/IR/Ops.cpp b/src/executable/mlir/Dialect/Python/IR/Ops.cpp index 3c63dcdc..e38bf0fc 100644 --- a/src/executable/mlir/Dialect/Python/IR/Ops.cpp +++ b/src/executable/mlir/Dialect/Python/IR/Ops.cpp @@ -406,7 +406,7 @@ namespace py { return region_or_block_arguments(getOperation(), successor); } - void BranchYieldOp::getSuccessorRegions(llvm::ArrayRef operands, + void BranchYieldOp::getSuccessorRegions(llvm::ArrayRef /*operands*/, llvm::SmallVectorImpl ®ions) { static_assert(BranchYieldOp::hasTrait< From f3cea1ee3a469765df61279e2bfe60f4ddf19fbc Mon Sep 17 00:00:00 2001 From: gf712 Date: Mon, 24 Aug 2026 09:34:28 +0100 Subject: [PATCH 12/12] mlir: fix transitive loop-else break/continue binding --- .../PythonToPythonBytecode.cpp | 67 ++++++++++++------- 1 file changed, 41 insertions(+), 26 deletions(-) diff --git a/src/executable/mlir/Conversion/PythonToPythonBytecode/PythonToPythonBytecode.cpp b/src/executable/mlir/Conversion/PythonToPythonBytecode/PythonToPythonBytecode.cpp index 3144bf6c..b3c4808d 100644 --- a/src/executable/mlir/Conversion/PythonToPythonBytecode/PythonToPythonBytecode.cpp +++ b/src/executable/mlir/Conversion/PythonToPythonBytecode/PythonToPythonBytecode.cpp @@ -56,39 +56,54 @@ namespace py { return mlir::isa(op); } - // True when `yield_op` is a loop-control (break/continue) yield that binds to - // the loop *enclosing* `loop` rather than to `loop` itself — i.e. it sits in - // `loop`'s orelse, which is not part of the loop body. - bool binds_to_enclosing_loop(mlir::py::PyLoopOpInterface loop, - mlir::py::BranchYieldOp yield_op) + // True when `yield_op`, a loop-control (break/continue) yield, binds to the + // loop whose body region is `body`. + // + // Python binds break/continue to the innermost loop whose *body* lexically + // contains it. An else clause is not part of its own loop's body, so a yield + // sitting there keeps searching outwards — and transitively so: an else nested + // inside another else is still lexically part of whatever body encloses the + // pair. Regions that are neither body nor orelse (a try body, a with body) are + // likewise transparent, which is what makes `break` inside a `try` bind to the + // loop around it. + bool binds_to_loop(mlir::Region &body, mlir::py::BranchYieldOp yield_op) { - return yield_op.getKind().has_value() && loop.isLoopOrelse(yield_op->getParentRegion()); + for (mlir::Region *region = yield_op->getParentRegion(); region != nullptr; + region = region->getParentRegion()) { + if (region == &body) { return true; } + auto loop = + mlir::dyn_cast_if_present(region->getParentOp()); + // A loop body (or a for's step) stops the search: the yield is that + // loop's, and its own pattern claims it. + if (loop && !loop.isLoopOrelse(region)) { return false; } + } + return false; } - // True when some loop nested in `region` still holds a break/continue that - // binds to the loop being lowered — i.e. one sitting in that nested loop's - // orelse. Such a yield cannot be rewritten yet: it lives in a region that has - // not been flattened, so branching it to our target block would be a - // cross-region block reference, which is invalid IR. + // True when a break/continue that binds to the loop whose body is `body` is + // somewhere replace_loop_branch_yields cannot reach yet — inside a nested + // region that has not been flattened into ours. Branching it to our target + // block now would be a cross-region block reference, which is invalid IR. // - // The caller defers (fails the match) until the nested loop lowers and inlines + // The caller defers (fails the match) until the nested op lowers and inlines // the yield into our region, the same innermost-first trick TryOpLowering uses - // for nested trys. Terminates because the innermost such loop has nothing - // nested to wait on. - bool has_pending_nested_orelse_control(mlir::Region ®ion) + // for nested trys. Terminates because the innermost such op has nothing nested + // to wait on. + bool has_pending_nested_orelse_control(mlir::Region &body) { - if (region.empty()) { return false; } + if (body.empty()) { return false; } bool pending = false; - region.walk([&pending](mlir::Operation *op) { - auto loop = mlir::dyn_cast(op); - if (!loop) { return WalkResult::advance(); } - loop.getLoopOrelseRegion().walk( - [&pending, loop](mlir::py::BranchYieldOp yield_op) { - if (binds_to_enclosing_loop(loop, yield_op)) { pending = true; } - }); - // Only this loop's own orelse matters here; anything deeper is the - // nested loop's problem and it defers on it in turn. - return WalkResult::skip(); + body.walk([&pending, &body](mlir::Operation *op) { + // Mirror replace_loop_branch_yields: what it walks through it rewrites + // in place, so only what it skips over can be pending. + if (!is_flattened_region_op(op)) { return WalkResult::advance(); } + op->walk([&pending, &body](mlir::py::BranchYieldOp yield_op) { + if (!yield_op.getKind().has_value()) { return WalkResult::advance(); } + if (!binds_to_loop(body, yield_op)) { return WalkResult::advance(); } + pending = true; + return WalkResult::interrupt(); + }); + return pending ? WalkResult::interrupt() : WalkResult::skip(); }); return pending; }