[mlir] Support DialectRegistry extension comparison #101119
Conversation
|
@llvm/pr-subscribers-flang-openmp @llvm/pr-subscribers-mlir Author: Nikhil Kalra (nikalra) ChangesThe current This change relaxes the assert such that the context dialect registry can be updated inside of a dynamic pass pipeline when in single-threaded mode. Full diff: https://github.com/llvm/llvm-project/pull/101119.diff 1 Files Affected:
diff --git a/mlir/lib/IR/MLIRContext.cpp b/mlir/lib/IR/MLIRContext.cpp
index 12336701c9ca0..14316874e4743 100644
--- a/mlir/lib/IR/MLIRContext.cpp
+++ b/mlir/lib/IR/MLIRContext.cpp
@@ -410,9 +410,10 @@ void MLIRContext::appendDialectRegistry(const DialectRegistry ®istry) {
if (registry.isSubsetOf(impl->dialectsRegistry))
return;
- assert(impl->multiThreadedExecutionContext == 0 &&
- "appending to the MLIRContext dialect registry while in a "
- "multi-threaded execution context");
+ assert(!impl->threadingIsEnabled ||
+ impl->multiThreadedExecutionContext == 0 &&
+ "appending to the MLIRContext dialect registry while in a "
+ "multi-threaded execution context");
registry.appendTo(impl->dialectsRegistry);
// For the already loaded dialects, apply any possible extensions immediately.
|
6b70652 to
0cf68e8
Compare
It seems like this should be a better approach by allowing us to check if extensions have already been registered into the context (instead of conservatively assuming they aren't). But, this change touches 90+ extra files, so if there's a way to do this more cleanly I'm more than happy to investigate that further! |
|
Thanks for giving it a try! I feel that relying on a string provided by the user is not desirable. We need to revisit this with the same kind of mechanism we used elsewhere (based on TypeID / global symbol address) somehow. You probably want to narrow this down to one example that we can iterate on before engaging in large scale changes. |
This reverts commit 26669bd.
| std::function<void(MLIRContext *, DialectsT * ...)>(extensionFn)); | ||
| } | ||
| template <typename... DialectsT> | ||
| void |
There was a problem hiding this comment.
We'll have to remove the std::function overload of addExtension so we can use the underlying function pointer as the TypeID, but the lost functionality can be achieved with a DialectExtension subclass if needed.
I think just a handful: we should get most of the lambdas for free since I don't think any of them have captures, but we'll have to add type ID macros to the DialectExtension subclasses. |
|
Can you add in the PR description a note about the removed API and the advice on how to proceed for people affected (how to replace it). |
Co-authored-by: Mehdi Amini <joker.eph@gmail.com>
Done! |
|
If this change looks good, would it be possible to merge it in? I don't have merge permissions for LLVM. Thanks! |
|
(Waiting for CI to complete) |
Actually this seems missing? Or maybe I wasn't clear: downstream folks who will broken by this patch would appreciate in the description a step-by-step list of actions they need to take to upgrade their codebase. |
Added step-by-step instructions for the two modified use cases: existing DialectExtensions and the removed overload. Please let me know if I'm missing anything there or if there's anything that won't be clear to affected users! |
|
Perfect, thanks! |
PassManager::runloads the dependent dialects for each pass into the current context prior to invoking the individual passes. If the dependent dialect is already loaded into the context, this should be a no-op. However, if there are extensions registered in theDialectRegistry, the dependent dialects are unconditionally registered into the context.This poses a problem for dynamic pass pipelines, however, because they will likely be executing while the context is in an immutable state (because of the parent pass pipeline being run).
To solve this, we'll update the extension registration API on
DialectRegistryto require a type ID for each extension that is registered. Then, instead of unconditionally registered dialects into a context if extensions are present, we'll check against the extension type IDs already present in the context's internalDialectRegistry. The context will only be marked as dirty if there are net-new extension types present in theDialectRegistrypopulated byPassManager::getDependentDialects.Note: this PR removes the
addExtensionoverload that utilizesstd::functionas the parameter. This is becausestd::functionis copyable and potentially allocates memory for the contained function so we can't use the function pointer as the unique type ID for the extension.Downstream changes required:
DialectExtensionsubclasses will need a type ID to be registered for each subclass. More details on how to register a type ID can be found here:llvm-project/mlir/include/mlir/Support/TypeID.h
Line 30 in 8b68e06
std::functionoverload ofaddExtensionwill need to be refactored into dedicatedDialectExtensionclasses with associated type IDs. The attachedstd::functioncan either be inlined into or called directly fromDialectExtension::apply.