[mlir] Add Dynamic Dialects

[math-fehr]

Summary

Dynamic dialects are dialects that can be defined at runtime.
Dynamic dialects are extensible by new operations, types, and
attributes at runtime.

[math-fehr]

I would like to add tests for dynamic dialects, by creating a new dynamic dialect, and then populating it with dynamic operations, types, and attributes.
However, the way mlir-opt currently works is by creating a DialectRegistry, and then passing it to MlirOptMain.

Is there any way of getting somewhere a MLIRContext to register the dialect?
Alternatively, how should we register DynamicDialect from DialectRegistry, since we want to directly access them to add operations/types/attributes?

[math-fehr]

Add parsing and printing of types/attributes

[math-fehr]

Add missing newline at end of file

[rriddle]

In D125201#3499771, @math-fehr wrote:

I would like to add tests for dynamic dialects, by creating a new dynamic dialect, and then populating it with dynamic operations, types, and attributes.
However, the way mlir-opt currently works is by creating a DialectRegistry, and then passing it to MlirOptMain.

Is there any way of getting somewhere a MLIRContext to register the dialect?
Alternatively, how should we register DynamicDialect from DialectRegistry, since we want to directly access them to add operations/types/attributes?

I think for the registry we could have a different type of allocator function that we wrap into the normal dialect allocator function. That way we could do
something like:

registry.insertDynamic("name", [](DynamicDialect *dialect) {
  // Insert the stuff for the dialect here.
});

Inside of the registry though, we should remove the need to create a TypeID upfront for dynamic dialects. Either by using a different registry map, or by using something like TypeID::get<void>() internally. Given that dynamic dialects don't have the same TypeID problems that normal dialects have, the TypeID in the registry doesn't really guard against much IMO.

mlir/include/mlir/IR/ExtensibleDialect.h

553

Is this explicit initialization necessary?

mlir/include/mlir/IR/MLIRContext.h

117–118

This should likely be named getOrLoadDynamicDialect to match the other API. With that, why do we need to abort if a dialect of the same name is already loaded?

mlir/lib/IR/ExtensibleDialect.cpp

549 ↗

(On Diff #427997)

Diagnostic messages should start with a lower case. Same above.

mlir/lib/IR/MLIRContext.cpp

461

If we are actually loading the dialect, we could copy the string into abstractDialectSymbolAllocator or something. That would remove a lot of the uniquing cost.

Also, can you assert that we aren't in a multithreaded context here?

[math-fehr]

Address comments, and add insertDynamic to DialectRegistry.

[math-fehr]

I added an insertDynamic in the DialectRegistry.

My dynamic.mlir test is currently failing (it seems because of a memory issue), and I'll debug this tomorrow.

mlir/include/mlir/IR/MLIRContext.h

117–118

I think I'll change it so it returns nullptr instead of aborting.

The problem with a getOrLoadDynamicDialect is that we wouldn't know if the dialect was created or not.
With getOrLoadDialect<MyDialect>, we do not care if the dialect was just loaded or not, since the resulting dialect will be the same. However, with a getOrLoadDynamicDialect, we would either get a dialect with already some operations/types/attributes defined, or a new dialect with nothing in it.

Does that make sense?

mlir/lib/IR/MLIRContext.cpp

461

Is this what you meant about abstractDialectSymbolAllocator?

Also, for the multithreading context, adding an assert here seems to fail on my test. Is there some place where I should deactivate it when loading a new dynamic dialect with the DialectRegistry during parsing?

[math-fehr]

Fix bug in dynamic dialect registration

[math-fehr]

In D125201#3505209, @math-fehr wrote:

I added an insertDynamic in the DialectRegistry.

My dynamic.mlir test is currently failing (it seems because of a memory issue), and I'll debug this tomorrow.

This is now fixed

[rriddle]

mlir/include/mlir/IR/MLIRContext.h

117–118

Can we just do what we do in the normal case? i.e. pass in the initializer function and only call it if the dialect needs to be constructed?

mlir/lib/IR/Dialect.cpp

175 ↗

(On Diff #428836)

Can we move nameStr here? or do something like nameStr = name.str() in the capture list?

mlir/lib/IR/MLIRContext.cpp

461

That feels weird, we can only load a dialect in a non-multithreaded context. We are already asserting in getOrLoadDialect, but we should also assert here as well (just to make sure the assert context is more accurate)

https://github.com/llvm/llvm-project/blob/cb778e9328292b7cfa4e0b43f26e4d1091923a39/mlir/lib/IR/MLIRContext.cpp#L424

mlir/test/lib/Dialect/TestDyn/TestDynDialect.cpp

19 ↗

(On Diff #428836)

Can we keep this in the test namespace? I don't think there is a good reason to diverge.

26 ↗

(On Diff #428836)

Could you move the functors out-of-line and the inline the call to DynamicOpDefinition::get?

29–31 ↗

(On Diff #428836)

mlir/test/lib/Dialect/TestDyn/TestDynDialect.h

1 ↗

(On Diff #428836)

This header shouldn't be necessary, it isn't included anywhere.

[rriddle]

Can you also add documentation on how to add/define a purely dynamic dialect?

[math-fehr]

Address comments

[math-fehr]

Remove unwanted line break

[math-fehr]

This should now be ready for review.
I added documentation for the dynamic dialect definition, and addressed the comments.

[rriddle]

Nice.

mlir/docs/DefiningDialects.md

396 ↗

(On Diff #431735)

mlir/include/mlir/IR/DialectRegistry.h

141–143 ↗

(On Diff #431735)

Can you reflow these comments?

mlir/include/mlir/IR/ExtensibleDialect.h

548–549

Can you reflow these comments?

mlir/include/mlir/IR/MLIRContext.h

114

117

Can we pass by pointer here? Dialects are generally interacted with by-pointer, so it would be nice to be consistent here.

mlir/lib/IR/Dialect.cpp

168–169 ↗

(On Diff #431735)

Can you reflow the comments here?

171 ↗

(On Diff #431735)

178 ↗

(On Diff #431735)

Is this to make sure the lambda returns the right type? Can you just mark the type in the lambda declaration? i.e. ...) -> Dialect * {

mlir/lib/IR/MLIRContext.cpp

468

The llvm:: shouldn't be necessary, can you drop it?

487

Can you just pass dialect by-value? It's a pointer so the reference doesn't matter.

986–987

This feels unrelated. Same with the comment changes above.

mlir/test/IR/dynamic.mlir

137 ↗

(On Diff #431735)

Super nit: can you drop the double newline?

143 ↗

(On Diff #431735)

Missing newline?

mlir/test/lib/Dialect/TestDyn/TestDynDialect.cpp

19–20 ↗

(On Diff #431735)

21 ↗

(On Diff #431735)

Can we drop the ::mlir:: on these? Same with the registry.

[math-fehr]

Address comments

[math-fehr]

I should have addressed the comments.
Could you land this for me?

[math-fehr]

Rebased on main

[math-fehr]

Sorry for the delay, I forgot a bit about this patch!
@rriddle, I just rebased it, could you land this for me, I do not have commit access?

[math-fehr]

Remove use of deprecated methods

[rriddle]

In D125201#3799952, @math-fehr wrote:

Remove use of deprecated methods

Is this fully ready? Will land now if so (oops).

[math-fehr]

Yes, this is ready! I just updated it to the main.

mlir/include/mlir/IR/ExtensibleDialect.h

private:
   friend DynamicOpDefinition;
 
   /// Allocates a type ID to uniquify operations.
   TypeID allocateTypeID() { return typeIDAllocator.allocate(); }
 
   /// Owns the TypeID generated at runtime for operations.
   TypeIDAllocator typeIDAllocator;
 };
+
+//===----------------------------------------------------------------------===//
+// Dynamic dialect
+//===----------------------------------------------------------------------===//
+
+/// A dialect that can be defined at runtime.
+/// It can be extended with new operations, types, and attributes at runtime.
+class DynamicDialect : public SelfOwningTypeID, public ExtensibleDialect {
+public:
+  DynamicDialect(StringRef name, MLIRContext *ctx)
+      : SelfOwningTypeID(),
+        ExtensibleDialect(name, ctx, SelfOwningTypeID::getTypeID()) {}
+
+  TypeID getTypeID() { return SelfOwningTypeID::getTypeID(); }
+};
 } // namespace mlir
 
 namespace llvm {
 /// Provide isa functionality for ExtensibleDialect.
 /// This is to override the isa functionality for Dialect.
 template <>
 struct isa_impl<mlir::ExtensibleDialect, mlir::Dialect> {
   static inline bool doit(const ::mlir::Dialect &dialect) {
     return mlir::ExtensibleDialect::classof(&dialect);
   }
 };
 } // namespace llvm
 
 #endif // MLIR_IR_EXTENSIBLEDIALECT_H

mlir/include/mlir/IR/MLIRContext.h

 class ThreadPool;
 } // namespace llvm
 
 namespace mlir {
 class DebugActionManager;
 class DiagnosticEngine;
 class Dialect;
 class DialectRegistry;
+class DynamicDialect;
 class InFlightDiagnostic;
 class Location;
 class MLIRContextImpl;
 class RegisteredOperationName;
 class StorageUniquer;
 
 /// MLIRContext is the top-level object for a collection of MLIR operations. It
 /// holds immortal uniqued objects like types, and the tables used to unique
public:
 
   /// Load a list dialects in the context.
   template <typename Dialect, typename OtherDialect, typename... MoreDialects>
   void loadDialect() {
     getOrLoadDialect<Dialect>();
     loadDialect<OtherDialect, MoreDialects...>();
   }
 
+  /// Create and load a new dialect that will be populated with operations,
+  /// types, and attributes at runtime. Abort if a dialect with a similar name
+  /// is already loaded.
+  DynamicDialect *createDynamicDialect(StringRef dialectNamespace);
+
   /// Load all dialects available in the registry in this context.
   void loadAllAvailableDialects();
 
   /// Get (or create) a dialect for the given derived dialect name.
   /// The dialect will be loaded from the registry if no dialect is found.
   /// If no dialect is loaded for this name and none is available in the
   /// registry, returns nullptr.
   Dialect *getOrLoadDialect(StringRef name);

mlir/lib/IR/MLIRContext.cpp

 #include "IntegerSetDetail.h"
 #include "TypeDetail.h"
 #include "mlir/IR/AffineExpr.h"
 #include "mlir/IR/AffineMap.h"
 #include "mlir/IR/Attributes.h"
 #include "mlir/IR/BuiltinDialect.h"
 #include "mlir/IR/Diagnostics.h"
 #include "mlir/IR/Dialect.h"
+#include "mlir/IR/ExtensibleDialect.h"
 #include "mlir/IR/IntegerSet.h"
 #include "mlir/IR/Location.h"
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/IR/Types.h"
 #include "mlir/Support/DebugAction.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/SetVector.h"
#endif
   std::unique_ptr<Dialect> &dialect = dialectIt->second;
   if (dialect->getTypeID() != dialectID)
     llvm::report_fatal_error("a dialect with namespace '" + dialectNamespace +
                              "' has already been registered");
 
   return dialect.get();
 }
 
+DynamicDialect *MLIRContext::createDynamicDialect(StringRef dialectNamespace) {
+  auto name = StringAttr::get(this, dialectNamespace);
+  auto *dialect = new DynamicDialect(name, this);
+  auto *loadedDialect =
+      getOrLoadDialect(name, dialect->getTypeID(), [&dialect]() {
+        return std::unique_ptr<DynamicDialect>(dialect);
+      });
+  // This is the same result as `getOrLoadDialect` (if it didn't failed),
+  // since it has the same TypeID, and TypeIDs are unique.
+  return dialect;
+}
+
 void MLIRContext::loadAllAvailableDialects() {
   for (StringRef name : getAvailableDialects())
     getOrLoadDialect(name);
 }
 
 llvm::hash_code MLIRContext::getRegistryHash() {
   llvm::hash_code hash(0);
   // Factor in number of loaded dialects, attributes, operations, types.
   hash = llvm::hash_combine(hash, impl->loadedDialects.size());
   hash = llvm::hash_combine(hash, impl->registeredAttributes.size());
   hash = llvm::hash_combine(hash, impl->registeredOperations.size());
   hash = llvm::hash_combine(hash, impl->registeredTypes.size());
   return hash;
 }
 
 bool MLIRContext::allowsUnregisteredDialects() {
   return impl->allowUnregisteredDialects;
 }
 
 void MLIRContext::allowUnregisteredDialects(bool allowing) {
   assert(impl->multiThreadedExecutionContext == 0 &&
          "changing MLIRContext `allow-unregistered-dialects` configuration "
          "while in a multi-threaded execution context");
   impl->allowUnregisteredDialects = allowing;
willBeValidAffineMap(unsigned dimCount, unsigned symbolCount,
   int64_t maxDimPosition = -1;
   int64_t maxSymbolPosition = -1;
   getMaxDimAndSymbol(ArrayRef<ArrayRef<AffineExpr>>(results), maxDimPosition,
                      maxSymbolPosition);
   if ((maxDimPosition >= dimCount) || (maxSymbolPosition >= symbolCount)) {
     LLVM_DEBUG(
         llvm::dbgs()
         << "maximum dimensional identifier position in result expression must "
            "be less than `dimCount` and maximum symbolic identifier position "
            "in result expression must be less than `symbolCount`\n");
     return false;
   }
   return true;
 }
 
 AffineMap AffineMap::get(MLIRContext *context) {
   return getImpl(/*dimCount=*/0, /*symbolCount=*/0, /*results=*/{}, context);
 }