Skip to content

Commit 05a5db5

Browse files
[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.
1 parent 2371d0a commit 05a5db5

File tree

5 files changed

+59
-13
lines changed

5 files changed

+59
-13
lines changed

mlir/include/mlir/Dialect/Transform/Transforms/Passes.td

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -51,8 +51,9 @@ def PreloadLibraryPass : Pass<"transform-preload-library"> {
5151

5252
Warning: Only a single such pass should exist for a given MLIR context.
5353
This is a temporary solution until a resource-based solution is available.
54-
TODO: investigate using a resource blob if some ownership mode allows it.
5554
}];
55+
// TODO: investigate using a resource blob if some ownership mode allows it.
56+
let dependentDialects = ["transform::TransformDialect"];
5657
let options = [
5758
ListOption<"transformLibraryPaths", "transform-library-paths", "std::string",
5859
"Optional paths to files with modules that should be merged into the "
@@ -67,6 +68,7 @@ def InterpreterPass : Pass<"transform-interpreter"> {
6768
sequence transformation specified by the provided name (defaults to
6869
`__transform_main`).
6970
}];
71+
let dependentDialects = ["transform::TransformDialect"];
7072
let options = [
7173
Option<"entryPoint", "entry-point", "std::string",
7274
/*default=*/[{"__transform_main"}],

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

Lines changed: 3 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -25,13 +25,10 @@ class InterpreterPass
2525
public:
2626
using Base::Base;
2727

28-
LogicalResult initialize(MLIRContext *context) override {
29-
// TODO: investigate using a resource blob if some ownership mode allows it.
30-
transformModule = transform::detail::getPreloadedTransformModule(context);
31-
return success();
32-
}
33-
3428
void runOnOperation() override {
29+
MLIRContext *context = &getContext();
30+
ModuleOp transformModule =
31+
transform::detail::getPreloadedTransformModule(context);
3532
if (failed(transform::applyTransformNamedSequence(
3633
getOperation(), transformModule,
3734
options.enableExpensiveChecks(true), entryPoint)))
@@ -41,11 +38,5 @@ class InterpreterPass
4138
private:
4239
/// Transform interpreter options.
4340
transform::TransformOptions options;
44-
45-
/// The separate transform module to be used for transformations, shared
46-
/// across multiple instances of the pass if it is applied in parallel to
47-
/// avoid potentially expensive cloning. MUST NOT be modified after the pass
48-
/// has been initialized.
49-
ModuleOp transformModule;
5041
};
5142
} // namespace
Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,17 @@
1+
// RUN: mlir-opt %s -transform-interpreter=entry-point=entry_point \
2+
// RUN: -split-input-file -verify-diagnostics
3+
4+
module attributes { transform.with_named_sequence } {
5+
transform.named_sequence @entry_point(!transform.any_op {transform.readonly}) {
6+
^bb0(%arg0: !transform.any_op):
7+
// expected-remark @below {{applying transformation}}
8+
transform.test_transform_op
9+
transform.yield
10+
}
11+
12+
transform.named_sequence @__transform_main(!transform.any_op {transform.readonly}) {
13+
^bb0(%arg0: !transform.any_op):
14+
transform.test_transform_op // Note: does not yield remark.
15+
transform.yield
16+
}
17+
}
Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,17 @@
1+
// RUN: mlir-opt %s -transform-interpreter \
2+
// RUN: -split-input-file -verify-diagnostics
3+
4+
module attributes { transform.with_named_sequence } {
5+
transform.named_sequence @__transform_main(!transform.any_op {transform.readonly}) {
6+
^bb0(%arg0: !transform.any_op):
7+
// expected-remark @below {{applying transformation}}
8+
transform.test_transform_op
9+
transform.yield
10+
}
11+
12+
transform.named_sequence @entry_point(!transform.any_op {transform.readonly}) {
13+
^bb0(%arg0: !transform.any_op):
14+
transform.test_transform_op // Note: does not yield remark.
15+
transform.yield
16+
}
17+
}
Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,19 @@
1+
// RUN: mlir-opt %s \
2+
// RUN: -transform-preload-library=transform-library-paths=%p%{fs-sep}test-interpreter-library \
3+
// RUN: -transform-interpreter=entry-point=private_helper \
4+
// RUN: -split-input-file -verify-diagnostics
5+
6+
// expected-remark @below {{message}}
7+
module {}
8+
9+
// -----
10+
11+
// Note: no remark here since local entry point takes precedence.
12+
module attributes { transform.with_named_sequence } {
13+
transform.named_sequence @private_helper(!transform.any_op {transform.readonly}) {
14+
^bb0(%arg0: !transform.any_op):
15+
// expected-remark @below {{applying transformation}}
16+
transform.test_transform_op
17+
transform.yield
18+
}
19+
}

0 commit comments

Comments
 (0)