This is an archive of the discontinued LLVM Phabricator instance.

Separate the Registration from Loading dialects in the Context
ClosedPublic

Authored by mehdi_amini on Aug 9 2020, 9:51 PM.

Details

Summary

This changes the behavior of constructing MLIRContext to no longer load globally registered dialects on construction. Instead Dialects are only loaded explicitly on demand:

  • the Parser is lazily loading Dialects in the context as it encounters them during parsing. This is the only purpose for registering dialects and not load them in the context.
  • Passes are expected to declare the dialects they will create entity from (Operations, Attributes, or Types), and the PassManager is loading Dialects into the Context when starting a pipeline.

This changes simplifies the configuration of the registration: a compiler only need to load the dialect for the IR it will emit, and the optimizer is self-contained and load the required Dialects. For example in the Toy tutorial, the compiler only needs to load the Toy dialect in the Context, all the others (linalg, affine, std, LLVM, ...) are automatically loaded depending on the optimization pipeline enabled.

Diff Detail

Event Timeline

mehdi_amini created this revision.Aug 9 2020, 9:51 PM
mehdi_amini requested review of this revision.Aug 9 2020, 9:51 PM

Passes are expected to declare the dialects they will create entity from (Operations, Attributes, or Types), and the PassManager is loading Dialects into the Context when starting a pipeline.

Why put this into PassManager? Why not let passes load lazily too only when needed and not do it up front for if it may be needed? What is the overhead of lazily loading and is the fear that it is too high? When is it needed? (E.g., does one need to have dialect loaded to create an op from that dialect? And if the input already has ops from that dialect it isn't really needed?).

Passes are expected to declare the dialects they will create entity from (Operations, Attributes, or Types), and the PassManager is loading Dialects into the Context when starting a pipeline.

Why put this into PassManager? Why not let passes load lazily too only when needed and not do it up front for if it may be needed? What is the overhead of lazily loading and is the fear that it is too high? When is it needed? (E.g., does one need to have dialect loaded to create an op from that dialect? And if the input already has ops from that dialect it isn't really needed?).

Excellent question! The reason is threading: front-loading this at the beginning of the pipeline is required to be lock-free.
We moved away from locking the context during registration, so these methods aren't thread-safe anymore (River had measurable perf improvements).

Looks good, did a first pass.

mlir/include/mlir/Conversion/Passes.td
69

nit: Wrap this at 80 characters.

79

nit: Move after constructor.

213

Move this before options.

299

Same here.

330

Same here and a few other places.

mlir/include/mlir/ExecutionEngine/JitRunner.h
24

Looks unused.

mlir/include/mlir/IR/Dialect.h
223

nit: allows for decoupling

227

The ordering of this container should be stable.

231

nit: add -> insert?

248

Can you change this pass in another dialect registry instead of the context? That would be more composable, and less confusing when looking at this compared to loadAll.

260

nit: What is the benefit of the non-const iterators here?

mlir/include/mlir/IR/MLIRContext.h
25

This is unnecessary.

rriddle added inline comments.Aug 13 2020, 12:48 PM
mlir/include/mlir/IR/MLIRContext.h
47

This seems like something the dialect registry would provide, why add this to the context?

139

This comment seems old.

mlir/include/mlir/Pass/Pass.h
63

nit: other than dialects that exist in the input. For example,

mlir/include/mlir/Pass/PassManager.h
112

nit: Add a comment here.

mlir/include/mlir/Transforms/Passes.td
105

This looks like a functional change?

mlir/lib/Conversion/PassDetail.h
18

nit: I would add that this is a forward declaration from DialectRegistry.h.

mlir/lib/Dialect/Affine/Transforms/PassDetail.h
16

Same here.

mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
1690

Can we autogen this with the dialect tablegen backend?

mlir/lib/IR/MLIRContext.cpp
451

nit: Use array_pod_sort here.

469

Why the extra ()?

mlir/lib/Parser/AttributeParser.cpp
250

nit: I would add a comment here to specifically call out that this is lazy loading a dialect.

mlir/lib/Parser/Parser.cpp
1473

nit: Check to see if the load was successful first. If the load fails, the second call is guaranteed to fail.

mlir/tools/mlir-opt/mlir-opt.cpp
178

nit: Drop trivial braces.

You could also use llvm::interleave here.

mehdi_amini marked 11 inline comments as done.

Address most comments

mehdi_amini marked 3 inline comments as done.

(remove spurious new line)

mehdi_amini marked 9 inline comments as done.Aug 13 2020, 10:47 PM
mehdi_amini added inline comments.
mlir/include/mlir/IR/MLIRContext.h
25

Leftover from previous iterations...

47

"Someone" needs to own the DialectRegistry: ultimately I'm trying to remove entirely the concept of global registry.

I originally didn't have a registry in the context, but that forces the caller of the parser to thread through separately the registry.
Instead we can register the dialects with the context (without loading them) and load them on-demand when parsing.

mlir/include/mlir/Transforms/Passes.td
105

Maybe a bad rebase? I don't remember changing this

mlir/lib/Conversion/PassDetail.h
18

From Dialect.h you mean?

mlir/lib/IR/MLIRContext.cpp
469

I can't read ternary well without these :)

I can remove it if it hurts your eyes

mehdi_amini marked 3 inline comments as done.

(revert unrelated change in mlir/Transforms/Passes.td)

Restore the change for bufferplacement: it is needed after all

Thanks for the review River! That was very much a draft...

mlir/include/mlir/Transforms/Passes.td
105

Actually the BufferPlacement pass was just not deriving from the generated base class but directly from FunctionPass, so this change is needed.

Update ctor initialization

mehdi_amini added inline comments.Aug 14 2020, 1:56 AM
mlir/include/mlir/IR/MLIRContext.h
42

@rriddle : FYI I added this ctor parameter for the transition. This ease landing this while clients have to update all their uses.

DavidTruby resigned from this revision.Aug 14 2020, 5:46 AM

Please fix the clang-tidy failures.

mlir/examples/toy/Ch5/mlir/LowerToAffineLoops.cpp
260

Can we update insert to be variadic?

mlir/include/mlir/Dialect/LLVMIR/LLVMOpBase.td
25 ↗(On Diff #285583)

nit: Drop the extra spaces surrounding "omp..."

mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
26 ↗(On Diff #285583)

Same here and a few other places.

mlir/include/mlir/IR/Dialect.h
228

You could also have used MapVector, but either is fine.

250

nit: name is in a different style.

257

Same here.

mlir/include/mlir/IR/MLIRContext.h
47

Is the registry on the context intended to be temporary then? i.e. will it be removed when the global registry is gone?

mlir/include/mlir/Pass/Pass.h
62

I think we should remove operations here, a dialect dependency can also arise from depending on attributes or types.

mlir/include/mlir/Pass/PassBase.td
81

Same here.

mlir/include/mlir/Support/MlirOptMain.h
26 ↗(On Diff #285583)

nit: an option upfront -> an upfront?

mlir/lib/CAPI/IR/IR.cpp
66

nit: Parameter comment on the false.

mlir/lib/Conversion/PassDetail.h
18

Yeah, I just guessed.

mlir/lib/IR/Dialect.cpp
34

nit: name.str()?

mlir/lib/IR/MLIRContext.cpp
451

Unresolved?

451

nit: Drop trivial braces here.

mlir/lib/Pass/Pass.cpp
747

nit: Fix the naming style here.

mlir/tools/mlir-tblgen/PassGen.cpp
108

nit: Naming style here.

flaub added a subscriber: flaub.Aug 14 2020, 5:21 PM
rriddle accepted this revision.Aug 14 2020, 5:34 PM
This revision is now accepted and ready to land.Aug 14 2020, 5:34 PM
mehdi_amini marked 15 inline comments as done.

Address comments

mehdi_amini added inline comments.Aug 14 2020, 11:50 PM
mlir/include/mlir/IR/Dialect.h
228

I thought about it, but MapVector depends on the order of insertion, and right now this is used with global initializer whose order isn't defined.
It would be deterministic across runs of the same binary (hopefully), but not across platform/environment or across revisions.

mlir/include/mlir/IR/MLIRContext.h
47

Addressed offline: it is only for allowing lazy loading in the parser

mlir/include/mlir/Pass/Pass.h
62
/// Register dependent dialects for the current pass.
/// A pass is expected to register the dialects it will create entities for
/// (Operations, Types, Attributes), other than dialect that exists in the
/// input. For example, a pass that converts from Linalg to Affine would
/// register the Affine dialect but does not need to register Linalg.
This revision was landed with ongoing or failed builds.Aug 15 2020, 1:08 AM
This revision was automatically updated to reflect the committed changes.
mehdi_amini reopened this revision.Aug 18 2020, 1:03 PM
This revision is now accepted and ready to land.Aug 18 2020, 1:03 PM
This revision was landed with ongoing or failed builds.Aug 18 2020, 4:25 PM
This revision was automatically updated to reflect the committed changes.

I fixed this earlier, tracking this build right now to see if flang tests
are passing:
http://lab.llvm.org:8011/builders/flang-aarch64-ubuntu/builds/2949

Build is back green!
http://lab.llvm.org:8011/builders/flang-aarch64-ubuntu/builds/2949

Sorry for the breakage: I tested with flang, but then I had a last minute
change of an API and only tested MLIR...
In general if a patch is breaking flang, feel free to revert directly and
notify in the revision!