[Mlir-commits] [mlir] [MLIR][Transform] Prefer entry points in current module (PR #151323)

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Erick Ochoa Lopez (amd-eochoalo)

<details>
<summary>Changes</summary>

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. 

---
Full diff: https://github.com/llvm/llvm-project/pull/151323.diff


2 Files Affected:

- (modified) mlir/lib/Dialect/Transform/Transforms/TransformInterpreterUtils.cpp (+76-10) 
- (added) mlir/test/Dialect/Transform/interpreter-entry-point-2.mlir (+24) 


``````````diff
diff --git a/mlir/lib/Dialect/Transform/Transforms/TransformInterpreterUtils.cpp b/mlir/lib/Dialect/Transform/Transforms/TransformInterpreterUtils.cpp
index 35ace1b2e0c3a..9ab484ff68078 100644
--- a/mlir/lib/Dialect/Transform/Transforms/TransformInterpreterUtils.cpp
+++ b/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;
   }
diff --git a/mlir/test/Dialect/Transform/interpreter-entry-point-2.mlir b/mlir/test/Dialect/Transform/interpreter-entry-point-2.mlir
new file mode 100644
index 0000000000000..e3e901a7eaf02
--- /dev/null
+++ b/mlir/test/Dialect/Transform/interpreter-entry-point-2.mlir
@@ -0,0 +1,24 @@
+// RUN: mlir-opt %s -transform-interpreter \
+// RUN:   -split-input-file -verify-diagnostics | 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
+  }
+}

``````````

</details>


https://github.com/llvm/llvm-project/pull/151323