Page MenuHomePhabricator

[mlir] Add Dynamic Dialects
AcceptedPublic

Authored by math-fehr on May 8 2022, 8:09 PM.

Details

Summary

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

Diff Detail

Event Timeline

math-fehr created this revision.May 8 2022, 8:09 PM
Herald added a project: Restricted Project. · View Herald Transcript
math-fehr requested review of this revision.May 8 2022, 8:09 PM

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 updated this revision to Diff 427995.May 9 2022, 1:08 AM

Add parsing and printing of types/attributes

math-fehr updated this revision to Diff 427997.May 9 2022, 1:11 AM

Add missing newline at end of file

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

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 marked 2 inline comments as done.

Address comments, and add insertDynamic to DialectRegistry.

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 updated this revision to Diff 428836.May 11 2022, 7:03 PM

Fix bug in dynamic dialect registration

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 added inline comments.May 12 2022, 11:20 PM
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
176

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
20

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

27

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

30–32
mlir/test/lib/Dialect/TestDyn/TestDynDialect.h
1 ↗(On Diff #428836)

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

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

math-fehr marked 7 inline comments as done.

Address comments

Remove unwanted line break

math-fehr marked 2 inline comments as done.May 24 2022, 11:19 AM

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

rriddle accepted this revision.May 24 2022, 10:10 PM

Nice.

mlir/docs/DefiningDialects.md
396
mlir/include/mlir/IR/DialectRegistry.h
141–150

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
169–170

Can you reflow the comments here?

172
179

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.

1009–1010

This feels unrelated. Same with the comment changes above.

mlir/test/IR/dynamic.mlir
137

Super nit: can you drop the double newline?

142

Missing newline?

mlir/test/lib/Dialect/TestDyn/TestDynDialect.cpp
20–21
22

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

This revision is now accepted and ready to land.May 24 2022, 10:10 PM
math-fehr updated this revision to Diff 431993.May 25 2022, 7:47 AM
math-fehr marked 15 inline comments as done.

Address comments

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