[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.

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"}],

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

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
+  }
+}

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
+  }
+}

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
+  }
+}