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

[llvmlistbot at llvm.org]

Author: Ingo Müller
Date: 2023-10-17T12:32:16+02:00
New Revision: 22e3bf4eaf6cbbd387a3789e7ee082434e62d072

URL: https://github.com/llvm/llvm-project/commit/22e3bf4eaf6cbbd387a3789e7ee082434e62d072
DIFF: https://github.com/llvm/llvm-project/commit/22e3bf4eaf6cbbd387a3789e7ee082434e62d072.diff

LOG: [mlir][transform] Fix new interpreter and library preloading passes. (#69190)

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.

Added: 
    mlir/test/Dialect/Transform/interpreter-entry-point.mlir
    mlir/test/Dialect/Transform/interpreter.mlir
    mlir/test/Dialect/Transform/preload-library.mlir

Modified: 
    mlir/include/mlir/Dialect/Transform/Transforms/Passes.td
    mlir/lib/Dialect/Transform/Transforms/InterpreterPass.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Dialect/Transform/Transforms/Passes.td b/mlir/include/mlir/Dialect/Transform/Transforms/Passes.td
index c900fee76b814d3..286f69bc52486cb 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 = ["::mlir::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 = ["::mlir::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..61d22252dc61dfd
--- /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
+  }
+}