[discussion] Where to register external models?

MLIR
1

External models are a way to implement an op interface in a place different from the op definition. E.g., in a different dialect/build unit/…

One unresolved question is where to register those external models.

Possible options are:

  • In the passes that require the external model. I.e., virtual void getDependentDialects(DialectRegistry &registry) const
  • When creating an MLIRContext. External models that ship with MLIR could be registered in registerAllDialects. People have already started registering their external models there recently.

The first option does not always work: If you have a pass that should operate on all ops that implement a certain external model, you may not be able to list all external model implementations. E.g., this is the case for bufferization: We want to have a generic --bufferize pass that bufferizes everything (including ops from user-defined dialects that do not ship with MLIR).

The second option seems like a good place to me. What are your thoughts?

[RFC] Dialect Extensions and "Promised" Interfaces
2

Thanks for bringing the topic up!
This is a question that I would abstract into “how to setup a project using MLIR” and there is a fundamental issue with external models that pokes at a hole in the MLIR infra right now.

We have always had two independent units of registrations: “dialect” and “pass”. Registration is even not needed at all for building a compiler: it is only useful when reading textual file (dialects registration) and building a pipeline from text (passes registration).

Setting up MLIR to use in a project is supposed to be “easy” because it is limited to picking the dialect you want to create IR for (the input dialect) and setup a pass pipeline programmatically. When doing so you don’t even ever need to register anything (assuming you don’t read MLIR textual file as input), you just load directly the dialect into the context (context->getOrLoadDialect<MyDialect>()) and you’re good to go.

The wrinkle is that this reliance on the dialects as the main unit of granularity for setting up MLIR does not account for things that don’t belong directly to a dialect. That is some canonicalization patterns that would match operations from multiple dialects, or similarly here the “external interfaces”.

I’m not sure what’s the right solution is for this, but I am quite strongly against using registerAllDialects for this and the existing two uses of registering external model should be removed from it: this isn’t intended to be in anyway a function that is part of the normal MLIR setup.

1 Like
3

This is a question that I would abstract into “how to setup a project using MLIR”

Every project has a program entry point. I don’t really know how people use MLIR in their own projects, but two examples that I have seen are mlir-opt and mlir-proto-opt. This seems like a good place to register everything. Kind of like a configuration file: You load everything that you need for your project in one central place and you’re good to go.

Registration is even not needed at all for building a compiler

Don’t you have to register a dialect before you can instantiate its operations?

I’m not sure what’s the right solution is for this, but I am quite strongly against using registerAllDialects for this and the existing two uses of registering external model should be removed from it.

I also think registerAllDialects is not really the right place because external models are not dialects. How about adding a new function registerAllExternalModels and calling that function in every place where we currently call registerAllDialects?

I think the purpose of registerAllDialects is:

  1. For testing: We have one binary mlir-opt that can run any test case. From any dialect. Basically it contains everything that MLIR has to offer.
  2. For lazy people: They are probably not using the entirety of MLIR but just a few select dialects/passes. So they should really be registering only those dialects/passes. But it’s easier to just register everything and call registerAllDialects. This is the case for the mlir-proto-opt mentioned above. Kind of like having dependencies in a BUILD file that are not needed; as long as it still compiles within reasonable time, they may not care.
4

No: you just need to load it in the Context when needed (ctx->getOrLoadDialect<TensorFlowDialect>() before emitting the TensorFlow operations.
More importantly: you don’t need to load or register anything else in the context, just setting up your pass pipeline is enough from there, thanks to getDependentDialects defined by the passes.

This would be a band-aid: any registerAllxxxx is only providing a convenient testing entry point and nothing more, I think you acknowledge this with your two bullet points.
So there needs to be an idiom there. We can take it by example: what if I want to write a compiler that goes from TOSA to LLVM?
Today I can load the TOSA dialect, emit operations, and then it is only about setting up the pass pipeline and nothing else. Now we need “more” somehow.

5

Now we need “more” somehow.

I’m trying to understand what we actually need here.

For dialects, it’s just a matter of loading them into the context. The MLIR user creates the context. So in a sense, users are responsible for loading the dialect. We don’t need registerAllDialects.

You could say the same about external models. The MLIR user could just register them at the same time when the context is created. If the external model is not needed in the project, there’s no need to register it.

But this could become tedious. As an example, there are many external models for bufferization and it would take many lines of code to register them in a project. Even harder to keep everything in sync when more models are added.

So we need some kind of functionality that automates the registration of external models (?). I’m not very familiar with MLIRContext etc., but maybe we could have some kind of helper function like registerAllExternalModels on the context. That function would iterate over all registered/loaded dialects and collect all external models that are defined within that dialect. (The dialects would have to expose them in some way.) For each external model, the function checks if the dialect of the operation that it extends is loaded. If so, the external model is registered; otherwise, it is not. registerAllExternalModels would be called once all dialects were loaded using getOrLoadDialect. Or maybe even after each call to getOrLoadDialect.

We would still have registerAllDialects and registerAllExternalModels in mlir-opt.cpp, but only for testing purposes.

6

Registration was intended to support dialect auto-loading while parsing textual IR. When constructing IR programmatically, it shouldn’t be necessary. Yet the external model API is provided on the DialectRegistry object that is used for registration, which ties external models to registration (rather than loading) unnecessarily. For the programmatic setup case, I suppose we could add something like MLIRContext::attachAttr/Op/TypeInterface that expects the dialect containing the attribute/op/type to have been already loaded. So the caller can issue context->attach*Interface() immediately after context->getOrLoadDialect().

Now this starts to break when we want to add an external model for the dialect that is not yet loaded because we are not creating objects from this IR, only some future pass will. I assume that we don’t want to preload all dialects that can be potentially emitted by passes (IIRC, this was the reason why dependentDialects exist in passes in the first place) just to register interfaces on them. Therefore, we need a mechanism to delay interface registration until the dialect is actually loaded (or not!). We still need to store the request for interface registration somewhere, and this happens to be DialectRegistry because we also want to handle such delayed interface registration requests when dialects are loaded by the parser.

At this point, I consider DialectRegistry to be not only a parser aid but a hook for things that must happen when the dialect is loaded, regardless of what triggered the loading. And I consider registration to be something that enables this hooks, not only parsing. We can consider splitting out the DialectPostLoadingHooks, but that seems to have little benefit as it would increase the overall system complexity and duplicate the information from the dialect registry (now one would have to “register” the dialect twice: once for parsing if desired and once for post-loading hooks).

Having registerAllExternalModels vs. a list of individual calls is a separate issue regarding convenience.

7

@ftynse : I see all of your explanation here as entirely correct, however I also think that this is missing the main problem. You are describing the mechanics of registration/loading and how to expose this to a client, while I am concerned mostly about the mental model associated with the abstract concepts at play here.
I am not concerned with “how” a user attaches an external model (DialectRegistry is fine…), but about how is a user supposed to even know about this.
@matthias-springer started to touch on this in this paragraph:

When assembling the compiler, I need not just to know about which dialect I’ll produce and which pipeline to assemble, but I also somehow need to know about a bunch of new extra things that will change the behavior of the compiler for a given pass pipeline.
I’m quite wary of the increase in complexity in all this right now, I’m looking into something similar to what I did with getDependentDialect: it allowed at the time to separate concerns and reduce maintenance burden (client don’t need to know about adding a dialect just because they added a new pass to the pipeline, or worse because a pass suddenly started to depend on this dialect).

1 Like
8

Here is another point that just came to mind: Code organization / build targets.

The external model implementations for bufferization are currently living in separate BUILD targets. E.g., the tensor dialect is living in MLIRTensor, but the external model impl for bufferization is living in MLIRTensorTransforms. This was a deliberate decision to keep the dependencies of dialect BUILD targets small.

Every bufferization external model impl depends at least on MLIRBufferization and MLIRMemRef (two other dialects). We probably don’t want every dialect that needs bufferization to depend on these two.

The implication of this design is that it is impossible to build any kind of “automatic registration” such as “register these external models automatically as soon as dialect X and dialect Y were registered”. This is not possible because the BUILD targets containing these external models may not even get linked in.

To address, @mehdi_amini’s concerns:

I am not concerned with “how” a user attaches an external model ( DialectRegistry is fine…), but about how is a user supposed to even know about this.

This is a good question and I have no solution for this. But we were having a very similar issue with the partial bufferization: How is a user supposed to know that they have to call -tensor-bufferize and -vector-bufferize and -finalizing-bufferize etc. Maybe the answer to this question is “documentation”.

In the case of bufferization, it could be something along the lines of:
When setting up your MLIRContext, you have to register all external model implementations of dialects that you wish to bufferize.

Same for other external models such as TensorInferTypeOpInterfaceImpl.

Ideally, there would be a way to automatically register external models just by linking in their respective BUILD targets. But I think that’s not possible.

9

All of this is about “pass pipeline setup” though: we kind of know how to abstract all this already. As I mentioned in my first answer:

We have always had two independent units of registrations: “dialect” and “pass”.

You need to know about the dialects you will emit directly, and you need to know how to build a pass pipeline, nothing more. This is a quite nicely decoupled system that is hard to get in an inconsistent state.

“Documentation” isn’t an acceptable answer to me here if this is about creating a system that is really easy to get in inconsistent and non-reproducible state.
Our current system has a high fidelity in terms of reproducer: given an an IR and a pass pipeline we should get consistent results.

10

Let’s take a look at the problem from this side: What is the purpose of external models? Why do we have them?

As far as I know, it’s about splitting code into small BUILD targets and keeping dependencies between BUILD targets small. Every op could just as well implement the interface directly instead of the external model. But that would introduce complex BUILD dependencies and may even result in cycles, meaning that we would sometimes have to merge multiple targets (e.g., two different dialects) into one big target.

To make external models as unsurprising for the user as possible, they should behave just as if the ops would implement their interfaces directly. Meaning: No external models registrations in getDependentDialects or any other mentioning of external models in a pass (or dialect).

Unfortunately, external models must be registered somewhere. Why not get done with it during startup and then we can forget about them. Like a BUILD file.

We have always had two independent units of registrations: “dialect” and “pass”.

The natural extension to me would be to have 3 independent units of registrations: “dialect”, “pass” and “external model”. No changes would be necessary to dialects or passes.

You need to know about the dialects you will emit directly, and you need to know how to build a pass pipeline, nothing more.

You will also have to know about external models now. In particular that they have to be registered. But this is much simpler than passes or dialects. It’s like having to know that you must add MLIRTensorTransforms to the LINK_LIBS section of your CMakeFiles.txt.

I still think during the setup of the MLIRContext would be a good place for registering external models. Along with the delayed registration mechanism that @ftynse described. As an idiom, every BUILD target that contains external models could have a public function that registers all the external models that it provides. Just like all those register(DialectName)Passes functions, except that they’re not generated through .td files but written manually.

Happy to do it in a different way, but it seems like nobody knows an alternative after a week of discussions.

Our current system has a high fidelity in terms of reproducer: given an an IR and a pass pipeline we should get consistent results.

I think this should not be affected. mlir-opt should have all external models registered during startup. If you’re talking about a user of MLIR that is not using mlir-opt, there is no guaranteed reproducibility anyway: When helping an external MLIR user, you may tell them to run pass -my-pass, but they may not even have -my-pass linked into their binary.

11

I don’t think there is a mechanism that allows the client not to know about external models in the general case. Otherwise, they would have been just regular interface implementations. So, for me, it is just a matter of minimizing the complexity of the API we provide.

Bufferization is a bit special in my opinion. It uses the external model mechanism (to reduce library dependencies), but the models not fully external, they are defined in the main codebase. In the fully external case, the external models would have been defined by the client that assembles the compiler, or at least taking from another codebase than the interface definition. In that case, the client already knows about these models to compile and link them, having to also register them does not sound that bad.

In bufferization (and the other interface as far as I can tell), external models are used to merely decouple dependencies of individual libraries. I would argue that this is the internal kitchen of MLIR and it is better off hidden from the client, which is what driven the inclusion of this registration into registerAllDialects.
What we are missing here is some concept for “details of pass X specific to dialect Y”, and potentially a mechanism to load those.

Let me elaborate. In the LLVM setup, there is a fixed set of instructions and passes know to which instructions they apply and how. In the current MLIR setup, there is no fixed set of operations so we flip the dependency direction: operations know about +/-all passes that apply to them and tell the passes how to apply via traits/interfaces. Both the former and the latter create many-to-one dependencies that may not be sustainable at scale. It looks like we want neither the passes to know about operations nor the operations to know about passes, so something else has to.

Not really. The original intent is to support out-of-tree interfaces. Imagine flang has some FlangTypeInfoInterface and wants some core MLIR dialects, e.g. builtin and LLVM, to implement this interface for their types. It is unreasonable to have MLIR depend on flang for this reason, and it is also unreasonable to have flang-specific interface included in MLIR so that core dialects can implement it. With external models, flang is able to provide both the interface and its implementations for the dialects that are external to it (hence the name). And at this point, it is perfectly reasonable IMO to require flang to register these models when assembling its flow.

Being able to split libraries within the same code base is arguably a side effect and we need to figure out whether this is something we want to encourage at all and, if so, do we reuse the same mechanism or amend it / split out the library idea.

1 Like
12

I would formulate this in a more abstract way than that, it is about layering and dependency injection. You allude to it by mentioning circular dependencies, it is just a symptom of layering issues that dependency injection allow to solve. Alex also illustrated another layering aspect when he elaborated on how dependency injection allows more decoupling (the example of Flang)

Well that’s the “non answer” to me, it does not address the points I raised above, it is just offloading all the complexity outside.

How so? This seems much more complex to me actually.

I don’t think so: build configuration are a trivial consequence from using an API. You need to link in only what you’re already using in the code, there is no “decision to make” in the build file that would affect the behavior.

Your two sentences are contradicting each other…
The behavior of the system will be affected by the fact that mlir-opt will register the external models, which will prevent from reproducing the issue of a system that is missing these.
Right now the worst thing that can happen is that mlir-opt wouldn’t be able to parse an input file if a dialect isn’t registered, but it would fail explicitly instead of just producing different result.

Another way to look at it: registering passes and dialect with mlir-opt has no effect until the user ask for a passes to be inserted in the pipeline or ask to parse an operation. Basically the registration is only ever touching the behavior of the user interaction, the provided inputs to mlir-opt, it won’t change the behavior of the system in any other way.

registerAllDialects being intended to exclusively be used for testing convenience, it can’t be a solution hiding anything from clients. I’m not even sure this API should really be public outside of MLIR.

Right, this is what motivated the external model in the first place I believe?

I agree with this, this is the internal kitchen of a particular compiler. That was the initial flow that we thought of when adding the external model (as far as I remember it).

We’re hitting composability issues now because we are trying to reuse this mechanism in the framework, outside of a given single compiler flow. Basically exposing external model and their registration “as a MLIR component”, that’s where it breaks down.

13

@mehdi_amini I heard from @ftynse that having different semantics around op casts could be another alternative of making users aware of the fact that external models must be registered. Could you elaborate a bit on that?

14

I brainstormed a bit about all this with @River707 last week, we should elaborate in a more concrete RFC. But in a nutshell the idea is centered around what is offered by the framework upstream where the constraints is that our “reusable components” can be assembled by users “easily” and avoid “error prone” patterns. So we could do:

  • A dialect upstream can declare that a given op implements an interface, without providing the implementation. For example we can declare that a given operation in the Tensor dialect implement the “bufferize interface”, but the dialect itself wouldn’t depend on the implementation.
  • a user of the Tensor dialect that won’t run the bufferization pass does not need to every provide / register an implementation for this interface.
  • a user running the “bufferization pass” on this operation, which will try to cast it to the “bufferize interface”, is in one of these two situations:
    • either the interface implementation has been registered (one way or another) and the cast succeeds.
    • or the implementation is missing (not registered) and cast will report_fatal_error in the compiler because of the broken promise here, which we would consider an incorrect setup of the framework.

The difference with the current state is that we don’t make this “promise” for implementing the interface in the dialect so we can’t check anything.

15

Sounds like a good plan.

We would still register all external models that ship with MLIR in registerAllDialects to make the test cases work, right?

16

TBD, but I’d rather have the registration handled directly in mlir-opt for example.

1 Like
17

Here is the proposal I was referring to earlier: [RFC] Dialect Extensions and "Promised" Interfaces

1 Like