[Mlir-commits] [mlir] [MLIR][Python] Support IsIsolatedFromAbove trait for python-defined operations (PR #211849)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Fri Jul 24 09:24:58 PDT 2026
https://github.com/PragmaTwice created https://github.com/llvm/llvm-project/pull/211849
This PR adds dynamic `IsIsolatedFromAbove` trait support and exposes it to Python as `IsIsolatedFromAboveTrait`.
The test verifies that a region with the trait cannot capture an SSA value from above, while a region without the trait can.
Assisted by Codex/GPT 5.6 Sol (for writing test cases).
>From 454f6a61f5032715da011cd59a469146ff18a02f Mon Sep 17 00:00:00 2001
From: PragmaTwice <twice at apache.org>
Date: Sat, 25 Jul 2026 00:15:16 +0800
Subject: [PATCH] [MLIR][Python] Support IsIsolatedFromAbove trait for
python-defined operations
---
mlir/include/mlir-c/ExtensibleDialect.h | 9 ++++
mlir/include/mlir/Bindings/Python/IRCore.h | 6 +++
mlir/include/mlir/IR/ExtensibleDialect.h | 8 +++
mlir/lib/Bindings/Python/IRCore.cpp | 23 ++++++++
mlir/lib/CAPI/IR/ExtensibleDialect.cpp | 8 +++
mlir/test/python/dialects/ext.py | 63 ++++++++++++++++++++++
6 files changed, 117 insertions(+)
diff --git a/mlir/include/mlir-c/ExtensibleDialect.h b/mlir/include/mlir-c/ExtensibleDialect.h
index c849658c710bb..f8dcb3f6800d3 100644
--- a/mlir/include/mlir-c/ExtensibleDialect.h
+++ b/mlir/include/mlir-c/ExtensibleDialect.h
@@ -54,6 +54,15 @@ mlirDynamicOpTraitIsTerminatorCreate(void);
/// terminator.
MLIR_CAPI_EXPORTED MlirTypeID mlirDynamicOpTraitIsTerminatorGetTypeID(void);
+/// Get the dynamic op trait that indicates regions are isolated from above.
+MLIR_CAPI_EXPORTED MlirDynamicOpTrait
+mlirDynamicOpTraitIsIsolatedFromAboveCreate(void);
+
+/// Get the type ID of the dynamic op trait that indicates regions are isolated
+/// from above.
+MLIR_CAPI_EXPORTED MlirTypeID
+mlirDynamicOpTraitIsIsolatedFromAboveGetTypeID(void);
+
/// Get the dynamic op trait that indicates regions have no terminator.
MLIR_CAPI_EXPORTED MlirDynamicOpTrait
mlirDynamicOpTraitNoTerminatorCreate(void);
diff --git a/mlir/include/mlir/Bindings/Python/IRCore.h b/mlir/include/mlir/Bindings/Python/IRCore.h
index 43a751bc9cd9f..3314e0b2a8fcf 100644
--- a/mlir/include/mlir/Bindings/Python/IRCore.h
+++ b/mlir/include/mlir/Bindings/Python/IRCore.h
@@ -1991,6 +1991,12 @@ class MLIR_PYTHON_API_EXPORTED NoTerminator : public PyDynamicOpTrait {
static void bind(nanobind::module_ &m);
};
+class MLIR_PYTHON_API_EXPORTED IsIsolatedFromAbove : public PyDynamicOpTrait {
+public:
+ static bool attach(const nanobind::object &opName, PyMlirContext &context);
+ static void bind(nanobind::module_ &m);
+};
+
} // namespace PyDynamicOpTraits
MLIR_PYTHON_API_EXPORTED MlirValue getUniqueResult(MlirOperation operation);
diff --git a/mlir/include/mlir/IR/ExtensibleDialect.h b/mlir/include/mlir/IR/ExtensibleDialect.h
index 15904358def11..2219d6fd57977 100644
--- a/mlir/include/mlir/IR/ExtensibleDialect.h
+++ b/mlir/include/mlir/IR/ExtensibleDialect.h
@@ -416,6 +416,14 @@ class IsTerminator : public DynamicOpTraitImpl<OpTrait::IsTerminator> {
class NoTerminator : public DynamicOpTraitImpl<OpTrait::NoTerminator> {};
+class IsIsolatedFromAbove
+ : public DynamicOpTraitImpl<OpTrait::IsIsolatedFromAbove> {
+public:
+ LogicalResult verifyRegionTrait(Operation *op) const override {
+ return OpTrait::impl::verifyIsIsolatedFromAbove(op);
+ }
+};
+
} // namespace DynamicOpTraits
//===----------------------------------------------------------------------===//
diff --git a/mlir/lib/Bindings/Python/IRCore.cpp b/mlir/lib/Bindings/Python/IRCore.cpp
index cceef55f9b15d..75cfd2a0a1c0b 100644
--- a/mlir/lib/Bindings/Python/IRCore.cpp
+++ b/mlir/lib/Bindings/Python/IRCore.cpp
@@ -2663,6 +2663,28 @@ void PyDynamicOpTraits::NoTerminator::bind(nb::module_ &m) {
nb::arg("op_name"), nb::arg("context").none() = nb::none());
}
+bool PyDynamicOpTraits::IsIsolatedFromAbove::attach(const nb::object &opName,
+ PyMlirContext &context) {
+ MlirDynamicOpTrait trait = mlirDynamicOpTraitIsIsolatedFromAboveCreate();
+ return attachOpTrait(opName, trait, context);
+}
+
+void PyDynamicOpTraits::IsIsolatedFromAbove::bind(nb::module_ &m) {
+ nb::class_<PyDynamicOpTraits::IsIsolatedFromAbove, PyDynamicOpTrait> cls(
+ m, "IsIsolatedFromAboveTrait");
+ cls.attr(typeIDAttr) =
+ PyTypeID(mlirDynamicOpTraitIsIsolatedFromAboveGetTypeID());
+ cls.attr("attach") = classmethod(
+ [](const nb::object &cls, const nb::object &opName,
+ DefaultingPyMlirContext context) {
+ return PyDynamicOpTraits::IsIsolatedFromAbove::attach(opName,
+ *context.get());
+ },
+ "Attach IsIsolatedFromAbove trait to the given operation name.",
+ nb::arg("cls"), nb::arg("op_name"),
+ nb::arg("context").none() = nb::none());
+}
+
} // namespace MLIR_BINDINGS_PYTHON_DOMAIN
} // namespace python
} // namespace mlir
@@ -5297,6 +5319,7 @@ void populateIRCore(nb::module_ &m) {
PyDynamicOpTrait::bind(m);
PyDynamicOpTraits::IsTerminator::bind(m);
PyDynamicOpTraits::NoTerminator::bind(m);
+ PyDynamicOpTraits::IsIsolatedFromAbove::bind(m);
// MLIRError exception.
MLIRError::bind(m);
diff --git a/mlir/lib/CAPI/IR/ExtensibleDialect.cpp b/mlir/lib/CAPI/IR/ExtensibleDialect.cpp
index 5c1f74884f2c5..2b4d5af0c9f13 100644
--- a/mlir/lib/CAPI/IR/ExtensibleDialect.cpp
+++ b/mlir/lib/CAPI/IR/ExtensibleDialect.cpp
@@ -54,6 +54,14 @@ MlirTypeID mlirDynamicOpTraitNoTerminatorGetTypeID() {
return wrap(DynamicOpTraits::NoTerminator::getStaticTypeID());
}
+MlirDynamicOpTrait mlirDynamicOpTraitIsIsolatedFromAboveCreate() {
+ return wrap(new DynamicOpTraits::IsIsolatedFromAbove());
+}
+
+MlirTypeID mlirDynamicOpTraitIsIsolatedFromAboveGetTypeID() {
+ return wrap(DynamicOpTraits::IsIsolatedFromAbove::getStaticTypeID());
+}
+
void mlirDynamicOpTraitDestroy(MlirDynamicOpTrait dynamicOpTrait) {
delete unwrap(dynamicOpTrait);
}
diff --git a/mlir/test/python/dialects/ext.py b/mlir/test/python/dialects/ext.py
index dde980df53dba..6e83da4a4a78a 100644
--- a/mlir/test/python/dialects/ext.py
+++ b/mlir/test/python/dialects/ext.py
@@ -637,6 +637,69 @@ class NoTermOp(TestRegion.Operation, name="no_term", traits=[NoTerminatorTrait])
print(e)
+# CHECK: TEST: testIsIsolatedFromAboveTrait
+ at run
+def testIsIsolatedFromAboveTrait():
+ class TestIsolated(Dialect, name="ext_isolated"):
+ pass
+
+ class UseOp(TestIsolated.Operation, name="use"):
+ value: Operand[Any]
+
+ class NotIsolatedOp(
+ TestIsolated.Operation, name="not_isolated", traits=[NoTerminatorTrait]
+ ):
+ body: Region
+
+ class IsolatedOp(
+ TestIsolated.Operation,
+ name="isolated",
+ traits=[NoTerminatorTrait, IsIsolatedFromAboveTrait],
+ ):
+ body: Region
+
+ with Context(), Location.unknown():
+ TestIsolated.load()
+
+ # CHECK: not isolated has trait: False
+ print(
+ "not isolated has trait:",
+ NotIsolatedOp.has_trait(IsIsolatedFromAboveTrait),
+ )
+ # CHECK: isolated has trait: True
+ print("isolated has trait:", IsolatedOp.has_trait(IsIsolatedFromAboveTrait))
+
+ not_isolated_module = Module.create()
+ with InsertionPoint(not_isolated_module.body):
+ i32 = IntegerType.get_signless(32)
+ value = arith.constant(i32, 0)
+ not_isolated = NotIsolatedOp()
+ not_isolated.body.blocks.append()
+ with InsertionPoint(not_isolated.body.blocks[0]):
+ UseOp(value)
+
+ assert not_isolated_module.operation.verify()
+ # CHECK: not isolated: verification succeeded
+ print("not isolated: verification succeeded")
+
+ isolated_module = Module.create()
+ with InsertionPoint(isolated_module.body):
+ value = arith.constant(i32, 0)
+ isolated = IsolatedOp()
+ isolated.body.blocks.append()
+ with InsertionPoint(isolated.body.blocks[0]):
+ UseOp(value)
+
+ try:
+ isolated_module.operation.verify()
+ except MLIRError as e:
+ # CHECK: isolated: Verification failed:
+ # CHECK: using value defined outside the region
+ print("isolated:", e)
+ else:
+ raise AssertionError("expected IsIsolatedFromAbove verification to fail")
+
+
# CHECK: TEST: testExtDialectWithType
@run
def testExtDialectWithType():
More information about the Mlir-commits
mailing list