[MLIR][Transform] Prefer entry points in current module (#151323)

The transform interpreter previously looked for the entry point using a
recursive walk in pre-order. This makes it so that any named_sequence
operation with an arbitrary level of nested-ness will be used as the
entry point for the transform interpreter as long as it is placed before
another one.

This change makes it so that code like the one reported in
https://github.com/llvm/llvm-project/issues/119578 works as expected.

Closes #119578 

Some comments: alternatively, it would also be possible to solve this
issue in a slightly more elegant manner. We could define a new walker
iterator that iterates through the operations in a breadth first search.

---------

Co-authored-by: Jakub Kuderski <kubakuderski@gmail.com>

mlir/lib/Dialect/Transform/Transforms/TransformInterpreterUtils.cpp

@@ -121,6 +121,80 @@ ModuleOp transform::detail::getPreloadedTransformModule(MLIRContext *context) {
       ->getLibraryModule();
 }
 
+static transform::TransformOpInterface
+findTransformEntryPointNonRecursive(Operation *op, StringRef entryPoint) {
+  for (Region &region : op->getRegions()) {
+    for (Block &block : region.getBlocks()) {
+      for (auto namedSequenceOp : block.getOps<transform::NamedSequenceOp>()) {
+        if (namedSequenceOp.getSymName() == entryPoint) {
+          return cast<transform::TransformOpInterface>(
+              namedSequenceOp.getOperation());
+        }
+      }
+    }
+  }
+  return nullptr;
+}
+
+static transform::TransformOpInterface
+findTransformEntryPointRecursive(Operation *op, StringRef entryPoint) {
+  transform::TransformOpInterface transform = nullptr;
+  op->walk<WalkOrder::PreOrder>(
+      [&](transform::NamedSequenceOp namedSequenceOp) {
+        if (namedSequenceOp.getSymName() == entryPoint) {
+          transform = cast<transform::TransformOpInterface>(
+              namedSequenceOp.getOperation());
+          return WalkResult::interrupt();
+        }
+        return WalkResult::advance();
+      });
+  return transform;
+}
+
+// Will look for the transform's entry point favouring NamedSequenceOps
+// ops that exist within the operation without the need for nesting.
+// If no operation exists in the blocks owned by op, then it will recursively
+// walk the op in preorder and find the first NamedSequenceOp that matches
+// the entry point's name.
+//
+// This allows for the following two use cases:
+// 1. op is a module annotated with the transform.with_named_sequence attribute
+//    that has an entry point in its block. E.g.,
+//
+//    ```mlir
+//    module {transform.with_named_sequence} {
+//      transform.named_sequence @__transform_main(%arg0 : !transform.any_op) ->
+//      () {
+//        transform.yield
+//      }
+//    }
+//    ```
+//
+// 2. op is a program which contains a nested module annotated with the
+//    transform.with_named_sequence attribute. E.g.,
+//
+//    ```mlir
+//    module {
+//      func.func @foo () {
+//      }
+//
+//      module {transform.with_named_sequence} {
+//        transform.named_sequence @__transform_main(%arg0 : !transform.any_op)
+//        -> () {
+//          transform.yield
+//        }
+//      }
+//    }
+//    ```
+static transform::TransformOpInterface
+findTransformEntryPointInOp(Operation *op, StringRef entryPoint) {
+  transform::TransformOpInterface transform =
+      findTransformEntryPointNonRecursive(op, entryPoint);
+  if (!transform)
+    transform = findTransformEntryPointRecursive(op, entryPoint);
+  return transform;
+}
+
 transform::TransformOpInterface
 transform::detail::findTransformEntryPoint(Operation *root, ModuleOp module,
                                            StringRef entryPoint) {
@@ -128,16 +202,8 @@ transform::detail::findTransformEntryPoint(Operation *root, ModuleOp module,
   if (module)
     l.push_back(module);
   for (Operation *op : l) {
-    transform::TransformOpInterface transform = nullptr;
-    op->walk<WalkOrder::PreOrder>(
-        [&](transform::NamedSequenceOp namedSequenceOp) {
-          if (namedSequenceOp.getSymName() == entryPoint) {
-            transform = cast<transform::TransformOpInterface>(
-                namedSequenceOp.getOperation());
-            return WalkResult::interrupt();
-          }
-          return WalkResult::advance();
-        });
+    TransformOpInterface transform =
+        findTransformEntryPointInOp(op, entryPoint);
     if (transform)
       return transform;
   }

mlir/test/Dialect/Transform/interpreter-entry-point-2.mlir

@@ -0,0 +1,23 @@
+// RUN: mlir-opt %s --transform-interpreter | FileCheck %s
+
+module @td_module_4 attributes {transform.with_named_sequence} {
+  module @foo_module attributes {transform.with_named_sequence} {
+    transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) -> () {
+      // CHECK: IR printer: foo_module top-level
+      transform.print {name="foo_module"}
+      transform.yield
+    }
+  }
+  module @bar_module attributes {transform.with_named_sequence} {
+    transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) -> () {
+      // CHECK: IR printer: bar_module top-level
+      transform.print {name="bar_module"}
+      transform.yield
+    }
+  }
+  transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) -> () {
+    transform.include @foo_module::@__transform_main failures(suppress) (%arg0) : (!transform.any_op) -> ()
+    transform.include @bar_module::@__transform_main failures(suppress) (%arg0) : (!transform.any_op) -> ()
+    transform.yield
+  }
+}