Page MenuHomePhabricator

[mlir] Add a hook for initializing passes before execution and use it in the Canonicalizer
ClosedPublic

Authored by rriddle on Dec 11 2020, 4:21 PM.

Details

Summary

This revision adds a new initialize(MLIRContext *) hook to passes that allows for them to initialize any heavy state before the first execution of the pass. A concrete use case of this is with patterns that rely on PDL, given that PDL is compiled at run time it is imperative that compilation results are cached as much as possible. The first use of this hook is in the Canonicalizer, which has the added benefit of reducing the number of expensive accesses to the context when collecting patterns.

Depends On D93146

Diff Detail

Event Timeline

rriddle created this revision.Dec 11 2020, 4:21 PM
rriddle requested review of this revision.Dec 11 2020, 4:21 PM

There is something subtle here: re-invoking the pass manager twice after having modified the context (like loading more dialect) may not reinitialize the passes correctly.
For example this fictional API:

std::pair<ModuleOp, ModuleOp> canonicalize_two_files(PassManager &PM, MLIRContext *ctx, StringRef file1, StringRef file2) {
ModuleOp m1 = parseMlirFile(ctx, file1);
PM.run(m1);
ModuleOp m2 = parseMlirFile(ctx, file2);
PM.run(m2);
return {m1, m2};
}

There is something subtle here: re-invoking the pass manager twice after having modified the context (like loading more dialect) may not reinitialize the passes correctly.
For example this fictional API:

std::pair<ModuleOp, ModuleOp> canonicalize_two_files(PassManager &PM, MLIRContext *ctx, StringRef file1, StringRef file2) {
ModuleOp m1 = parseMlirFile(ctx, file1);
PM.run(m1);
ModuleOp m2 = parseMlirFile(ctx, file2);
PM.run(m2);
return {m1, m2};
}

Yes. The main reason to cache the initialization is to properly handle the dynamic pipeline case. For the static pipeline, we can guarantee that it is only done once per full execution. For dynamic pipelines, we have no such guarantee or knowledge that it was initialized. We could enforce that the user initializes the pipeline.

There is something subtle here: re-invoking the pass manager twice after having modified the context (like loading more dialect) may not reinitialize the passes correctly.
For example this fictional API:

std::pair<ModuleOp, ModuleOp> canonicalize_two_files(PassManager &PM, MLIRContext *ctx, StringRef file1, StringRef file2) {
ModuleOp m1 = parseMlirFile(ctx, file1);
PM.run(m1);
ModuleOp m2 = parseMlirFile(ctx, file2);
PM.run(m2);
return {m1, m2};
}

Yes. The main reason to cache the initialization is to properly handle the dynamic pipeline case. For the static pipeline, we can guarantee that it is only done once per full execution. For dynamic pipelines, we have no such guarantee or knowledge that it was initialized. We could enforce that the user initializes the pipeline.

I don't particularly like the caching, but pushed to phab so we can discuss it here.

rriddle updated this revision to Diff 314676.Jan 5 2021, 11:38 AM

Change to using "Generations" for initialization, and add docs

There is something subtle here: re-invoking the pass manager twice after having modified the context (like loading more dialect) may not reinitialize the passes correctly.
For example this fictional API:

std::pair<ModuleOp, ModuleOp> canonicalize_two_files(PassManager &PM, MLIRContext *ctx, StringRef file1, StringRef file2) {
ModuleOp m1 = parseMlirFile(ctx, file1);
PM.run(m1);
ModuleOp m2 = parseMlirFile(ctx, file2);
PM.run(m2);
return {m1, m2};
}

Yes. The main reason to cache the initialization is to properly handle the dynamic pipeline case. For the static pipeline, we can guarantee that it is only done once per full execution. For dynamic pipelines, we have no such guarantee or knowledge that it was initialized. We could enforce that the user initializes the pipeline.

I don't particularly like the caching, but pushed to phab so we can discuss it here.

Changed to using "generations", which removes all additional state from Pass. This also allows for pass managers to be reinitialization, but avoids unnecessary reinitialization for dynamic pass pipelines.

mehdi_amini accepted this revision.Jan 6 2021, 2:37 PM
mehdi_amini added inline comments.
mlir/docs/PassManagement.md
147

?

This revision is now accepted and ready to land.Jan 6 2021, 2:37 PM
rriddle updated this revision to Diff 315501.Jan 8 2021, 1:24 PM
rriddle marked an inline comment as done.

Resolve comments and rebase

This revision was landed with ongoing or failed builds.Jan 8 2021, 1:38 PM
This revision was automatically updated to reflect the committed changes.