[Mlir-commits] [mlir] 81bdba1 - [MLIR] Add HasAncestor op trait (#195447)

[llvmlistbot at llvm.org]

Author: Henrich Lauko
Date: 2026-05-05T07:00:30+02:00
New Revision: 81bdba1fee08fc81ca0f0a2fa07a72acfd4f001a

URL: https://github.com/llvm/llvm-project/commit/81bdba1fee08fc81ca0f0a2fa07a72acfd4f001a
DIFF: https://github.com/llvm/llvm-project/commit/81bdba1fee08fc81ca0f0a2fa07a72acfd4f001a.diff

LOG: [MLIR] Add HasAncestor op trait (#195447)

Add HasAncestor/AncestorOneOf traits that verify an operation has a
specific ancestor anywhere in the parent chain, unlike HasParent which
only checks the immediate parent.

Added: 
    

Modified: 
    mlir/docs/Traits/_index.md
    mlir/include/mlir/IR/OpBase.td
    mlir/include/mlir/IR/OpDefinition.h
    mlir/test/IR/traits.mlir
    mlir/test/lib/Dialect/Test/TestOps.td

Removed: 
    


################################################################################
diff  --git a/mlir/docs/Traits/_index.md b/mlir/docs/Traits/_index.md
index 4127fe75dc6f5..866716b9f5193 100644
--- a/mlir/docs/Traits/_index.md
+++ b/mlir/docs/Traits/_index.md
@@ -272,6 +272,16 @@ trait. In particular, broadcasting behavior is not allowed. See the comments on
 This trait provides APIs and verifiers for operations that can only be nested
 within regions that are attached to operations of `ParentOpType`.
 
+### HasAncestor
+
+*   `OpTrait::HasAncestor<typename AncestorOpType>` -- `HasAncestor<string op>`
+    or `AncestorOneOf<list<string> opList>`
+
+This trait provides APIs and verifiers for operations that must appear somewhere
+inside a region attached to an operation of `AncestorOpType`. Unlike `HasParent`,
+which checks only the immediate parent, `HasAncestor` walks the full ancestor
+chain.
+
 ### IsolatedFromAbove
 
 *   `OpTrait::IsIsolatedFromAbove` -- `IsolatedFromAbove`

diff  --git a/mlir/include/mlir/IR/OpBase.td b/mlir/include/mlir/IR/OpBase.td
index 7f36e6c74c7f7..1e34959d0d557 100644
--- a/mlir/include/mlir/IR/OpBase.td
+++ b/mlir/include/mlir/IR/OpBase.td
@@ -141,6 +141,14 @@ class ParentOneOf<list<string> ops>
     : ParamNativeOpTrait<"HasParent", !interleave(ops, ", ")>,
       StructuralOpTrait;
 
+// Op has the provided one as an ancestor (anywhere in the parent chain).
+class HasAncestor<string op> : ParamNativeOpTrait<"HasAncestor", op>,
+                               StructuralOpTrait;
+
+class AncestorOneOf<list<string> ops>
+    : ParamNativeOpTrait<"HasAncestor", !interleave(ops, ", ")>,
+      StructuralOpTrait;
+
 // Op result type is derived from the first attribute. If the attribute is an
 // subclass of `TypeAttrBase`, its value is used, otherwise, the type of the
 // attribute content is used.

diff  --git a/mlir/include/mlir/IR/OpDefinition.h b/mlir/include/mlir/IR/OpDefinition.h
index e886ac45675e2..c1fba10e06a90 100644
--- a/mlir/include/mlir/IR/OpDefinition.h
+++ b/mlir/include/mlir/IR/OpDefinition.h
@@ -1324,6 +1324,32 @@ struct HasParent {
   };
 };
 
+/// This class provides a verifier for ops that are expecting an ancestor
+/// (anywhere up the parent chain) to be one of the given op types.
+template <typename... AncestorOpTypes>
+struct HasAncestor {
+  template <typename ConcreteType>
+  class Impl : public TraitBase<ConcreteType, Impl> {
+  public:
+    static LogicalResult verifyTrait(Operation *op) {
+      if (op->getParentOfType<AncestorOpTypes...>())
+        return success();
+
+      return op->emitOpError()
+             << "expects ancestor op "
+             << (sizeof...(AncestorOpTypes) != 1 ? "to be one of '" : "'")
+             << llvm::ArrayRef({AncestorOpTypes::getOperationName()...}) << "'";
+    }
+
+    template <typename AncestorOpType =
+                  std::tuple_element_t<0, std::tuple<AncestorOpTypes...>>>
+    std::enable_if_t<sizeof...(AncestorOpTypes) == 1, AncestorOpType>
+    getAncestorOp() {
+      return this->getOperation()->template getParentOfType<AncestorOpType>();
+    }
+  };
+};
+
 /// A trait for operations that have an attribute specifying operand segments.
 ///
 /// Certain operations can have multiple variadic operands and their size

diff  --git a/mlir/test/IR/traits.mlir b/mlir/test/IR/traits.mlir
index 49cfd7e496746..b20edde07a6c7 100644
--- a/mlir/test/IR/traits.mlir
+++ b/mlir/test/IR/traits.mlir
@@ -303,6 +303,85 @@ func.func @failedParentOneOf_wrong_parent1() {
    }) : () -> ()
 }
 
+// -----
+
+// CHECK: succeededHasAncestor_direct_parent
+func.func @succeededHasAncestor_direct_parent() {
+  "test.parent"() ({
+    "test.child_with_ancestor"() : () -> ()
+    "test.finish"() : () -> ()
+  }) : () -> ()
+  return
+}
+
+// -----
+
+// CHECK: succeededHasAncestor_indirect
+func.func @succeededHasAncestor_indirect() {
+  "test.parent"() ({
+    "some.intermediate_op"() ({
+      "test.child_with_ancestor"() : () -> ()
+      "test.finish"() : () -> ()
+    }) : () -> ()
+    "test.finish"() : () -> ()
+  }) : () -> ()
+  return
+}
+
+// -----
+
+func.func @failedHasAncestor_wrong_ancestor() {
+  "some.op"() ({
+    // expected-error at +1 {{'test.child_with_ancestor' op expects ancestor op 'test.parent'}}
+    "test.child_with_ancestor"() : () -> ()
+  }) : () -> ()
+}
+
+// -----
+
+// CHECK: succeededAncestorOneOf
+func.func @succeededAncestorOneOf() {
+  "test.parent"() ({
+    "test.child_with_ancestor_one_of"() : () -> ()
+    "test.finish"() : () -> ()
+  }) : () -> ()
+  return
+}
+
+// -----
+
+// CHECK: succeededAncestor1OneOf
+func.func @succeededAncestor1OneOf() {
+  "test.parent1"() ({
+    "test.child_with_ancestor_one_of"() : () -> ()
+    "test.finish"() : () -> ()
+  }) : () -> ()
+  return
+}
+
+// -----
+
+// CHECK: succeededAncestorOneOf_indirect
+func.func @succeededAncestorOneOf_indirect() {
+  "test.parent1"() ({
+    "some.intermediate_op"() ({
+      "test.child_with_ancestor_one_of"() : () -> ()
+      "test.finish"() : () -> ()
+    }) : () -> ()
+    "test.finish"() : () -> ()
+  }) : () -> ()
+  return
+}
+
+// -----
+
+func.func @failedAncestorOneOf_wrong_ancestor() {
+  "some.otherop"() ({
+    // expected-error at +1 {{'test.child_with_ancestor_one_of' op expects ancestor op to be one of 'test.parent, test.parent1'}}
+    "test.child_with_ancestor_one_of"() : () -> ()
+    "test.finish"() : () -> ()
+  }) : () -> ()
+}
 
 // -----
 

diff  --git a/mlir/test/lib/Dialect/Test/TestOps.td b/mlir/test/lib/Dialect/Test/TestOps.td
index 348ff5d7f4ea0..b5a6cd535d9f2 100644
--- a/mlir/test/lib/Dialect/Test/TestOps.td
+++ b/mlir/test/lib/Dialect/Test/TestOps.td
@@ -890,6 +890,15 @@ def ParentOp1 : TEST_Op<"parent1"> {
 def ChildWithParentOneOf : TEST_Op<"child_with_parent_one_of",
                                 [ParentOneOf<["ParentOp", "ParentOp1"]>]>;
 
+// HasAncestor trait
+def ChildWithAncestor
+    : TEST_Op<"child_with_ancestor", [HasAncestor<"ParentOp">]>;
+
+// AncestorOneOf trait
+def ChildWithAncestorOneOf
+    : TEST_Op<"child_with_ancestor_one_of", [AncestorOneOf<["ParentOp",
+                                                            "ParentOp1"]>]>;
+
 def TerminatorOp : TEST_Op<"finish", [Terminator]>;
 def SingleBlockImplicitTerminatorOp : TEST_Op<"SingleBlockImplicitTerminator",
     [SingleBlockImplicitTerminator<"TerminatorOp">]> {