[Mlir-commits] [mlir] [mlir][transform] Fix new interpreter and library preloading passes. (PR #69190)

[Ingo Müller]

https://github.com/ingomueller-net updated https://github.com/llvm/llvm-project/pull/69190

>From 05a5db5748a30915ce0ad51aa893af7396aff0aa Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ingo=20M=C3=BCller?= <ingomueller at google.com>
Date: Mon, 16 Oct 2023 10:52:20 +0000
Subject: [PATCH] [mlir][transform] Fix new interpreter and library preloading
 passes.

This PR fixes the two recently added passes from #68661, which were
non-functional and untested. In particular:
* The passes did not declare their dependent dialects, so they could not
  run at all in the most simple cases.
* The mechanism of loading the library module in the initialization of
  the intepreter pass is broken by design (but, fortunately, also not
  necessary). This is because the initialization of all passes happens
  before the execution of any other pass, so the "preload library" pass
  has not run yet at the time the interpreter pass gets initialized.
  Instead, the library is now loaded every time the interpreter pass is
  run. This should not be exceedingly expensive, since it only consists
  of looking up the library in the dialect. Also, this removes the
  library module from the pass state, making it possible in the future
  to preload libraries in several passes.
* The PR adds tests for the two passes, which were completely untested
  previously.
---
 .../Dialect/Transform/Transforms/Passes.td    |  4 +++-
 .../Transform/Transforms/InterpreterPass.cpp  | 15 +++------------
 .../Transform/interpreter-entry-point.mlir    | 17 +++++++++++++++++
 mlir/test/Dialect/Transform/interpreter.mlir  | 17 +++++++++++++++++
 .../Dialect/Transform/preload-library.mlir    | 19 +++++++++++++++++++
 5 files changed, 59 insertions(+), 13 deletions(-)
 create mode 100644 mlir/test/Dialect/Transform/interpreter-entry-point.mlir
 create mode 100644 mlir/test/Dialect/Transform/interpreter.mlir
 create mode 100644 mlir/test/Dialect/Transform/preload-library.mlir

diff --git a/mlir/include/mlir/Dialect/Transform/Transforms/Passes.td b/mlir/include/mlir/Dialect/Transform/Transforms/Passes.td
index c900fee76b814d3..8d53bf175f906a6 100644
--- a/mlir/include/mlir/Dialect/Transform/Transforms/Passes.td
+++ b/mlir/include/mlir/Dialect/Transform/Transforms/Passes.td
@@ -51,8 +51,9 @@ def PreloadLibraryPass : Pass<"transform-preload-library"> {
 
     Warning: Only a single such pass should exist for a given MLIR context.
     This is a temporary solution until a resource-based solution is available.
-    TODO: investigate using a resource blob if some ownership mode allows it.
   }];
+  // TODO: investigate using a resource blob if some ownership mode allows it.
+  let dependentDialects = ["transform::TransformDialect"];
   let options = [
     ListOption<"transformLibraryPaths", "transform-library-paths", "std::string",
     "Optional paths to files with modules that should be merged into the "
@@ -67,6 +68,7 @@ def InterpreterPass : Pass<"transform-interpreter"> {
     sequence transformation specified by the provided name (defaults to
     `__transform_main`).
   }];
+  let dependentDialects = ["transform::TransformDialect"];
   let options = [
     Option<"entryPoint", "entry-point", "std::string",
            /*default=*/[{"__transform_main"}],
diff --git a/mlir/lib/Dialect/Transform/Transforms/InterpreterPass.cpp b/mlir/lib/Dialect/Transform/Transforms/InterpreterPass.cpp
index f473d5aa728c519..3ec51d88729a0e7 100644
--- a/mlir/lib/Dialect/Transform/Transforms/InterpreterPass.cpp
+++ b/mlir/lib/Dialect/Transform/Transforms/InterpreterPass.cpp
@@ -25,13 +25,10 @@ class InterpreterPass
 public:
   using Base::Base;
 
-  LogicalResult initialize(MLIRContext *context) override {
-    // TODO: investigate using a resource blob if some ownership mode allows it.
-    transformModule = transform::detail::getPreloadedTransformModule(context);
-    return success();
-  }
-
   void runOnOperation() override {
+    MLIRContext *context = &getContext();
+    ModuleOp transformModule =
+        transform::detail::getPreloadedTransformModule(context);
     if (failed(transform::applyTransformNamedSequence(
             getOperation(), transformModule,
             options.enableExpensiveChecks(true), entryPoint)))
@@ -41,11 +38,5 @@ class InterpreterPass
 private:
   /// Transform interpreter options.
   transform::TransformOptions options;
-
-  /// The separate transform module to be used for transformations, shared
-  /// across multiple instances of the pass if it is applied in parallel to
-  /// avoid potentially expensive cloning. MUST NOT be modified after the pass
-  /// has been initialized.
-  ModuleOp transformModule;
 };
 } // namespace
diff --git a/mlir/test/Dialect/Transform/interpreter-entry-point.mlir b/mlir/test/Dialect/Transform/interpreter-entry-point.mlir
new file mode 100644
index 000000000000000..ccd9bef3d506ddf
--- /dev/null
+++ b/mlir/test/Dialect/Transform/interpreter-entry-point.mlir
@@ -0,0 +1,17 @@
+// RUN: mlir-opt %s -transform-interpreter=entry-point=entry_point \
+// RUN:   -split-input-file -verify-diagnostics
+
+module attributes { transform.with_named_sequence } {
+  transform.named_sequence @entry_point(!transform.any_op {transform.readonly}) {
+  ^bb0(%arg0: !transform.any_op):
+    // expected-remark @below {{applying transformation}}
+    transform.test_transform_op
+    transform.yield
+  }
+
+  transform.named_sequence @__transform_main(!transform.any_op {transform.readonly}) {
+  ^bb0(%arg0: !transform.any_op):
+    transform.test_transform_op // Note: does not yield remark.
+    transform.yield
+  }
+}
diff --git a/mlir/test/Dialect/Transform/interpreter.mlir b/mlir/test/Dialect/Transform/interpreter.mlir
new file mode 100644
index 000000000000000..bb41420bef4d63c
--- /dev/null
+++ b/mlir/test/Dialect/Transform/interpreter.mlir
@@ -0,0 +1,17 @@
+// RUN: mlir-opt %s -transform-interpreter \
+// RUN:   -split-input-file -verify-diagnostics
+
+module attributes { transform.with_named_sequence } {
+  transform.named_sequence @__transform_main(!transform.any_op {transform.readonly}) {
+  ^bb0(%arg0: !transform.any_op):
+    // expected-remark @below {{applying transformation}}
+    transform.test_transform_op
+    transform.yield
+  }
+
+  transform.named_sequence @entry_point(!transform.any_op {transform.readonly}) {
+  ^bb0(%arg0: !transform.any_op):
+    transform.test_transform_op // Note: does not yield remark.
+    transform.yield
+  }
+}
diff --git a/mlir/test/Dialect/Transform/preload-library.mlir b/mlir/test/Dialect/Transform/preload-library.mlir
new file mode 100644
index 000000000000000..55ae51bcd63b4b4
--- /dev/null
+++ b/mlir/test/Dialect/Transform/preload-library.mlir
@@ -0,0 +1,19 @@
+// RUN: mlir-opt %s \
+// RUN:   -transform-preload-library=transform-library-paths=%p%{fs-sep}test-interpreter-library \
+// RUN:   -transform-interpreter=entry-point=private_helper \
+// RUN:   -split-input-file -verify-diagnostics
+
+// expected-remark @below {{message}}
+module {}
+
+// -----
+
+// Note: no remark here since local entry point takes precedence.
+module attributes { transform.with_named_sequence } {
+  transform.named_sequence @private_helper(!transform.any_op {transform.readonly}) {
+  ^bb0(%arg0: !transform.any_op):
+    // expected-remark @below {{applying transformation}}
+    transform.test_transform_op
+    transform.yield
+  }
+}
\ No newline at end of file