[RFC] Dialect Extensions and "Promised" Interfaces

MLIR
#1

The current state of dialect extensibility is kind of awkward and clunky, and the current system is not really scalable or sustainable in general. This RFC seeks to start addressing this by introducing general mechanisms to help solve some of the current problems that we are facing.

Background

There are various ways that our current mechanisms are unsustainable, but the main way that this manifests is due to the fact that we shove all dialect add-ons (such as interfaces) into a single dialect library. This results in a bloated dialect library with lots of conditionally necessary dependencies that result from trying to package everything together. A recent discussion on where to register external models is a good example of some of the problems and frustrations arising from this. Another simple in-tree example that can help illustrate this is the Inliner. The inliner transformation defines a dialect interface to be used to specify how to inline a dialect and when it’s legal. For most dialects this is innocuous, but the inliner support for the standard dialect may create cf(ControlFlow) operations as part of its implementation. This means that the standard dialect needs to add a dependency on the ControlFlow dialect, as its inliner interface may end up creating ControlFlow operations. This may seem fairly simple, but this dependency is only necessary for the inlining transformation. If a user doesn’t need the inliner (which does happen mind you), they are still required to take on this additional dependency. If you apply this to other compilation flow specific transformations (such as bufferization), the point of unnecessary and untenable bloat becomes more quickly apparent.

Do we have anything that can currently be done to fix that?

Well, the current way of extending a dialect (or dialect owned construct, e.g. Attribute/Op/Type) is by adding a delayed interface to the DialectRegistry. This can handle simple cases, but it has several drawbacks: we can’t add conditional dialect dependencies, it only handles interfaces and each interface type requires a specific API, there is no control over when the interface gets added, etc. Another fundamental problem with our current system for extending a dialect is that we don’t have any indications on misconfiguration. One of the nice things about having everything in a single dialect library is that we know that the interfaces will always be loaded. We don’t have to worry about a user remembering to load the inliner support for the standard dialect, because we packaged it with the dialect itself.

Proposal

This proposal is split into two main parts:

DialectExtension

I propose we scrap the current delayed interface support from the DialectRegistry, and instead add a more general DialectExtension construct. This construct is essentially a callback that is invoked when a dialect (or set of dialects) have been loaded. This greatly simplifies the API surface area on the registry, provides a more convenient grouping mechanism for related add-ons, and also opens the door for more powerful dialect extensions. For example, if we take the standard dialect inliner example from before, we could now express inliner support as an extension:

/// This extension is applied when the `StandardOpsDialect` is loaded.
void mlir::standard::registerInlinerExtension(DialectRegistry &registry) {
  registry.addExtension(+[](MLIRContext *ctx, StandardOpsDialect *dialect) {
    dialect->addInterfaces<StdInlinerInterface>();

    // The inliner extension relies on the ControlFlow dialect.
    ctx->getOrLoadDialect<cf::ControlFlowDialect>();
  });
}

The above is an example using a simple callback, but the underlying mechanism is a new DialectExtension class:

template <typename DerivedT, typename... DialectsT>
class DialectExtension : public DialectExtensionBase {
public:
  /// Applies this extension to the given context and set of required dialects.
  virtual void apply(MLIRContext *context, DialectsT *...dialects) const = 0;
};

/// We can define the Inliner extension above as:
class InlinerExtension : public DialectExtension<InlinerExtension, StandardOpsDialect> {
public:
  void apply(MLIRContext *context, StandardOpsDialect *dialect) const override {
    dialect->addInterfaces<StdInlinerInterface>();

    // The inliner extension relies on the ControlFlow dialect.
    ctx->getOrLoadDialect<cf::ControlFlowDialect>();
  };
};

void mlir::standard::registerInlinerExtension(DialectRegistry &registry) {
 registry.addExtensions<InlinerExtension>();
}

Another example of something that this could also open up for the future, is using extensions to register canonicalization patterns involving multiple dialects without those dialects needing to explicitly know about each other.

Promised Interfaces

Above introduced the concept of a DialectExtension, but what that infra doesn’t solve is one of the final points in the background section:

One of the nice things about having everything in a single dialect library is that we know that the
interfaces will always be loaded. We don't have to worry about a user remembering to load the inliner
support for the standard dialect, because we packaged it with the dialect itself.

If we exposed the Standard dialect inliner interface as an extension, we get the benefit of a more scalable dialect library/reduced dependencies/etc., but we don’t effectively guard the system against misconfiguration. To that end, I also propose that we introduce the concept of a “Promised” Interface. A “promised” interface is essentially an interface that a dialect (or its constructs, e.g. attributes/ops/types/etc.) asserts that it has an implementation for. What this boils down to is that a dialect will claim it supports an interface (which could be a DialectInterface/AttrInterface/OpInterface/TypeInterface/etc.) …

void StandardOpsDialect::initialize() {
  ...
  declarePromisedInterface<DialectInlinerInterface>();
}

… with the expectation that the interface is loaded via an extension. If at any point the interface is properly attached, such as when applying an extension, the promise is resolved. If the interface is never attached and the interface is attempted to be used (e.g. via cast/isa/etc.) we can inform the user of the misconfiguration:

checking for an interface (`DialectInlinerInterface`) that was promised by dialect 'std' but never 
implemented. This is generally an indication that the dialect extension implementing the interface was 
never registered.

Final Thoughts

Our current system isn’t maintainable or scalable. Needing to shove every analysis and transformation dependency into the main dialect library is untenable, and often not possible (e.g. it could introduce circular dependencies). We need to start developing a system with which we can scale analysis and transformation additions to dialects in a clean and composable way. This does mean that more things will need to be registered, but the intention is to build out this infrastructure in such a way that it is harder to get wrong (e.g. users shouldn’t be left wondering why a transformation suddenly doesn’t work when an extension was not registered properly). As with any of the registration based infra, this is an ever evolving process and I don’t think we will start at the perfect end state.

I’ve uploaded two proof of concept patches at ⚙ D120367 [mlir] Refactor DialectRegistry delayed interface support into a general DialectExtension mechanism and ⚙ D120368 [mlir] Add support for "promised" interfaces

– River

4 Likes
[discussion] Where to register external models?
LLVM Weekly - #426, February 28th 2022
[doc] mlir-translate / mlir-opt
#2

+1 for the DialectExtension part!

I also like the idea of promised interfaces. This should give users clear instructions on what to do if they forget to register an external model. And avoid surprising behavior. Just one question:

Does this work for op interfaces that are defined in other dialects? E.g., BufferizableOpInterface is in Dialect/Bufferization/IR/BufferizableOpInterface.td, so it seems like every dialect that would promise an interface implementation would need to depend on the Bufferization dialect.

#3

Yes, and this applies to the Inliner interface example as well (which is technically defined in TransformUtils). The only thing that is important is the type name of the interface (for the declarePromisedInterface API), which can be grabbed (e.g. via forward declaration) in ways such that you don’t depend on the library where the interface is defined.

– River

#4

Also the interface is co-located under dialect/Bufferization but can it be in its own library that the dialect depends on instead of being in the same library?

#5

Just to make sure we are on the same page. I would say for promised interfaces in particular, we do not want the dialect library depending on the interface library. The extension library of course, but not the dialect itself.

Otherwise yeah, interfaces should be in a focus library.

– River

#6

This sounds great to me – to get rid of a dialect’s dependencies on other dialects where those dependencies were actually only needed for transformation/analysis purposes as opposed to for the purposes of defining, folding, or canonicalizing ops. As pointed out, even for the latter part (canonicalization), this may be a way to avoid inter-dialect library dependencies but one could argue that the dependency is the right thing there instead of magically/surprisingly rewriting/canonicalizing ops depending on whether extensions are loaded.

I didn’t understand the question here. The whole point here is to free the dialect from depending on another dialect for the purposes of interface implementation. Why would you have the dialect depend on the “dialect extension” library?! It’s the other way around: the registerInlinerExtension would depend on both the dialects involved. Both, the interface and the dialect extension would have to be in a separate library and depend on the requisite dialects.

A minor question here. Where/In which library does this code live? Would it say be a combined StdDialectExtensions with all dialect interface implementations in it? I don’t think there is a need to create one library per interface?!

#7

For extensions, it’s really a packaging question. If we always package all extensions together, we don’t really avoid the dependency issues we had before; we kind of just shift it around. We still want end users to be able to manage necessary dependencies in some way. Of course, it would be equally annoying to register every extension for every dialect individually. IMO right now, we will likely end up grouping extensions around different components/features/dependency sets. Bufferization is a great motivating example of the need for separation here, because it is very specific to a certain compilation flow. Unless you are originating from a tensor based compiler (and even then you might not want to use the bufferization functionality in-tree), you likely never want to depend on the bufferization interfaces/dialect/etc. With that being said, I still foresee us providing all encompassing StdDialectExtensions-like libraries for tools/bootstrapping/users who just don’t care about dependencies initially. These “all-extension” libraries provide a nice starting point for maintaining the current state, with the ability to carve out more constrained extension sets as we move forward.

As an aside, I don’t have an idea of what the best state is here. We will almost certainly build additional functionality on top of and evolve what is being proposed here (as we did with dialect dependencies/registration). Once we have the bottom layer and start trying to use it more in practice, we will likely more easily spot integration pain points and develop better layers on top.

– River

#8

To be able to declare the promise in the dialect you need to depend on “something” though? For example the “BufferizationInterface”. There has to be some sort of dependency to achieve this as far as I can tell.
So the Tensor dialect may still have to depend on the bufferization interface to declare the promise, even though the Tensor dialect would not depend on the bufferization dialect, only the implementation of the interface for the Tensor dialect would.