[mlir] Method to iterate over registered operations for a given dialect class. #112344
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Member
|
@llvm/pr-subscribers-mlir-core Author: Rajveer Singh Bharadwaj (Rajveer100) ChangesPart of #111591 Currently we have Full diff: https://github.com/llvm/llvm-project/pull/112344.diff 3 Files Affected:
diff --git a/mlir/include/mlir/IR/MLIRContext.h b/mlir/include/mlir/IR/MLIRContext.h
index d17bbac81655b5..cfad6874b8f4a9 100644
--- a/mlir/include/mlir/IR/MLIRContext.h
+++ b/mlir/include/mlir/IR/MLIRContext.h
@@ -195,7 +195,12 @@ class MLIRContext {
/// Return a sorted array containing the information about all registered
/// operations.
- ArrayRef<RegisteredOperationName> getRegisteredOperations();
+ SmallVector<RegisteredOperationName, 0> getRegisteredOperations();
+
+ /// Return a sorted array containing the information for registered operations
+ /// filtered by dialect name.
+ SmallVector<RegisteredOperationName, 0>
+ getRegisteredOperationsByDialect(StringRef dialectName);
/// Return true if this operation name is registered in this context.
bool isOperationRegistered(StringRef name);
diff --git a/mlir/lib/IR/MLIRContext.cpp b/mlir/lib/IR/MLIRContext.cpp
index f05666fcde207b..bb0da94e985517 100644
--- a/mlir/lib/IR/MLIRContext.cpp
+++ b/mlir/lib/IR/MLIRContext.cpp
@@ -188,7 +188,11 @@ class MLIRContextImpl {
/// This is a sorted container of registered operations for a deterministic
/// and efficient `getRegisteredOperations` implementation.
- SmallVector<RegisteredOperationName, 0> sortedRegisteredOperations;
+ SmallVector<std::pair<StringRef, RegisteredOperationName>, 0>
+ sortedRegisteredOperations;
+
+ /// This returns the number of registered operations for a given dialect.
+ llvm::DenseMap<StringRef, size_t> getCountByDialectName;
/// This is a list of dialects that are created referring to this context.
/// The MLIRContext owns the objects. These need to be declared after the
@@ -707,8 +711,31 @@ void MLIRContext::printStackTraceOnDiagnostic(bool enable) {
}
/// Return information about all registered operations.
-ArrayRef<RegisteredOperationName> MLIRContext::getRegisteredOperations() {
- return impl->sortedRegisteredOperations;
+SmallVector<RegisteredOperationName, 0> MLIRContext::getRegisteredOperations() {
+ SmallVector<RegisteredOperationName, 0> operations;
+ std::transform(impl->sortedRegisteredOperations.begin(),
+ impl->sortedRegisteredOperations.end(),
+ std::back_inserter(operations),
+ [](const auto &t) { return t.second; });
+
+ return operations;
+}
+
+/// Return information for registered operations by dialect.
+SmallVector<RegisteredOperationName, 0>
+MLIRContext::getRegisteredOperationsByDialect(StringRef dialectName) {
+ SmallVector<RegisteredOperationName, 0> operations;
+
+ auto lowerBound = std::lower_bound(
+ impl->sortedRegisteredOperations.begin(),
+ impl->sortedRegisteredOperations.end(), std::make_pair(dialectName, ""),
+ [](auto &lhs, auto &rhs) { return lhs.first.compare(rhs.first); });
+ auto count = impl->getCountByDialectName[dialectName];
+
+ std::transform(lowerBound, lowerBound + count, std::back_inserter(operations),
+ [](const auto &t) { return t.second; });
+
+ return operations;
}
bool MLIRContext::isOperationRegistered(StringRef name) {
@@ -976,14 +1003,20 @@ void RegisteredOperationName::insert(
"operation name registration must be successful");
// Add emplaced operation name to the sorted operations container.
- RegisteredOperationName &value = emplaced.first->second;
- ctxImpl.sortedRegisteredOperations.insert(
- llvm::upper_bound(ctxImpl.sortedRegisteredOperations, value,
- [](auto &lhs, auto &rhs) {
- return lhs.getIdentifier().compare(
- rhs.getIdentifier());
- }),
- value);
+ StringRef dialectClass = impl->getDialect()->getNamespace();
+ ctxImpl.getCountByDialectName[dialectClass] += 1;
+
+ std::pair<StringRef, RegisteredOperationName> value = {
+ dialectClass, emplaced.first->second};
+
+ auto upperBound = llvm::upper_bound(
+ ctxImpl.sortedRegisteredOperations, value, [](auto &lhs, auto &rhs) {
+ if (lhs.first == rhs.first)
+ return lhs.second.getIdentifier().compare(rhs.second.getIdentifier());
+ return lhs.first.compare(rhs.first);
+ });
+
+ ctxImpl.sortedRegisteredOperations.insert(upperBound, value);
}
//===----------------------------------------------------------------------===//
diff --git a/mlir/lib/Rewrite/FrozenRewritePatternSet.cpp b/mlir/lib/Rewrite/FrozenRewritePatternSet.cpp
index 17fe02df9f66cd..d3317fc6d4fe30 100644
--- a/mlir/lib/Rewrite/FrozenRewritePatternSet.cpp
+++ b/mlir/lib/Rewrite/FrozenRewritePatternSet.cpp
@@ -73,7 +73,7 @@ FrozenRewritePatternSet::FrozenRewritePatternSet(
// Functor used to walk all of the operations registered in the context. This
// is useful for patterns that get applied to multiple operations, such as
// interface and trait based patterns.
- std::vector<RegisteredOperationName> opInfos;
+ SmallVector<RegisteredOperationName> opInfos;
auto addToOpsWhen =
[&](std::unique_ptr<RewritePattern> &pattern,
function_ref<bool(RegisteredOperationName)> callbackFn) {
|
Member
|
@llvm/pr-subscribers-mlir Author: Rajveer Singh Bharadwaj (Rajveer100) ChangesPart of #111591 Currently we have Full diff: https://github.com/llvm/llvm-project/pull/112344.diff 3 Files Affected:
diff --git a/mlir/include/mlir/IR/MLIRContext.h b/mlir/include/mlir/IR/MLIRContext.h
index d17bbac81655b5..cfad6874b8f4a9 100644
--- a/mlir/include/mlir/IR/MLIRContext.h
+++ b/mlir/include/mlir/IR/MLIRContext.h
@@ -195,7 +195,12 @@ class MLIRContext {
/// Return a sorted array containing the information about all registered
/// operations.
- ArrayRef<RegisteredOperationName> getRegisteredOperations();
+ SmallVector<RegisteredOperationName, 0> getRegisteredOperations();
+
+ /// Return a sorted array containing the information for registered operations
+ /// filtered by dialect name.
+ SmallVector<RegisteredOperationName, 0>
+ getRegisteredOperationsByDialect(StringRef dialectName);
/// Return true if this operation name is registered in this context.
bool isOperationRegistered(StringRef name);
diff --git a/mlir/lib/IR/MLIRContext.cpp b/mlir/lib/IR/MLIRContext.cpp
index f05666fcde207b..bb0da94e985517 100644
--- a/mlir/lib/IR/MLIRContext.cpp
+++ b/mlir/lib/IR/MLIRContext.cpp
@@ -188,7 +188,11 @@ class MLIRContextImpl {
/// This is a sorted container of registered operations for a deterministic
/// and efficient `getRegisteredOperations` implementation.
- SmallVector<RegisteredOperationName, 0> sortedRegisteredOperations;
+ SmallVector<std::pair<StringRef, RegisteredOperationName>, 0>
+ sortedRegisteredOperations;
+
+ /// This returns the number of registered operations for a given dialect.
+ llvm::DenseMap<StringRef, size_t> getCountByDialectName;
/// This is a list of dialects that are created referring to this context.
/// The MLIRContext owns the objects. These need to be declared after the
@@ -707,8 +711,31 @@ void MLIRContext::printStackTraceOnDiagnostic(bool enable) {
}
/// Return information about all registered operations.
-ArrayRef<RegisteredOperationName> MLIRContext::getRegisteredOperations() {
- return impl->sortedRegisteredOperations;
+SmallVector<RegisteredOperationName, 0> MLIRContext::getRegisteredOperations() {
+ SmallVector<RegisteredOperationName, 0> operations;
+ std::transform(impl->sortedRegisteredOperations.begin(),
+ impl->sortedRegisteredOperations.end(),
+ std::back_inserter(operations),
+ [](const auto &t) { return t.second; });
+
+ return operations;
+}
+
+/// Return information for registered operations by dialect.
+SmallVector<RegisteredOperationName, 0>
+MLIRContext::getRegisteredOperationsByDialect(StringRef dialectName) {
+ SmallVector<RegisteredOperationName, 0> operations;
+
+ auto lowerBound = std::lower_bound(
+ impl->sortedRegisteredOperations.begin(),
+ impl->sortedRegisteredOperations.end(), std::make_pair(dialectName, ""),
+ [](auto &lhs, auto &rhs) { return lhs.first.compare(rhs.first); });
+ auto count = impl->getCountByDialectName[dialectName];
+
+ std::transform(lowerBound, lowerBound + count, std::back_inserter(operations),
+ [](const auto &t) { return t.second; });
+
+ return operations;
}
bool MLIRContext::isOperationRegistered(StringRef name) {
@@ -976,14 +1003,20 @@ void RegisteredOperationName::insert(
"operation name registration must be successful");
// Add emplaced operation name to the sorted operations container.
- RegisteredOperationName &value = emplaced.first->second;
- ctxImpl.sortedRegisteredOperations.insert(
- llvm::upper_bound(ctxImpl.sortedRegisteredOperations, value,
- [](auto &lhs, auto &rhs) {
- return lhs.getIdentifier().compare(
- rhs.getIdentifier());
- }),
- value);
+ StringRef dialectClass = impl->getDialect()->getNamespace();
+ ctxImpl.getCountByDialectName[dialectClass] += 1;
+
+ std::pair<StringRef, RegisteredOperationName> value = {
+ dialectClass, emplaced.first->second};
+
+ auto upperBound = llvm::upper_bound(
+ ctxImpl.sortedRegisteredOperations, value, [](auto &lhs, auto &rhs) {
+ if (lhs.first == rhs.first)
+ return lhs.second.getIdentifier().compare(rhs.second.getIdentifier());
+ return lhs.first.compare(rhs.first);
+ });
+
+ ctxImpl.sortedRegisteredOperations.insert(upperBound, value);
}
//===----------------------------------------------------------------------===//
diff --git a/mlir/lib/Rewrite/FrozenRewritePatternSet.cpp b/mlir/lib/Rewrite/FrozenRewritePatternSet.cpp
index 17fe02df9f66cd..d3317fc6d4fe30 100644
--- a/mlir/lib/Rewrite/FrozenRewritePatternSet.cpp
+++ b/mlir/lib/Rewrite/FrozenRewritePatternSet.cpp
@@ -73,7 +73,7 @@ FrozenRewritePatternSet::FrozenRewritePatternSet(
// Functor used to walk all of the operations registered in the context. This
// is useful for patterns that get applied to multiple operations, such as
// interface and trait based patterns.
- std::vector<RegisteredOperationName> opInfos;
+ SmallVector<RegisteredOperationName> opInfos;
auto addToOpsWhen =
[&](std::unique_ptr<RewritePattern> &pattern,
function_ref<bool(RegisteredOperationName)> callbackFn) {
|
Member
Author
Member
Author
|
Before proceeding with additional test cases, let me know if this works well. I am sure there are many logical ways to approach this. |
joker-eph
reviewed
Oct 15, 2024
joker-eph
reviewed
Oct 15, 2024
Contributor
|
Per usual, Mehdi beat me to the review. Thank you for the PR. |
Rajveer100
force-pushed
the
subrange-dialect-ops
branch
from
a63e19f to
41ec7af
Compare
River707
requested changes
Oct 16, 2024
Rajveer100
force-pushed
the
subrange-dialect-ops
branch
from
41ec7af to
859a8ac
Compare
Member
Author
|
Let me know if the updated logic works well. |
![graphite-app[bot]](https://avatars.githubusercontent.com/in/158384?s=60&v=4){kind=small}
Rajveer100
force-pushed
the
subrange-dialect-ops
branch
from
859a8ac to
0ca27ec
Compare
Rajveer100
force-pushed
the
subrange-dialect-ops
branch
from
0ca27ec to
30677fb
Compare
joker-eph
reviewed
Oct 16, 2024
Rajveer100
force-pushed
the
subrange-dialect-ops
branch
from
30677fb to
469021b
Compare
Rajveer100
force-pushed
the
subrange-dialect-ops
branch
2 times, most recently
from
a0849d7 to
3c94c22
Compare
Member
Author
|
I have made the changes, I think it should make sense now. PS: This could have been way smoother :\ |
River707
reviewed
Oct 16, 2024
Rajveer100
force-pushed
the
subrange-dialect-ops
branch
from
3c94c22 to
644ab88
Compare
Rajveer100
force-pushed
the
subrange-dialect-ops
branch
from
644ab88 to
d123ad6
Compare
joker-eph
reviewed
Oct 16, 2024
joker-eph
reviewed
Oct 16, 2024
Rajveer100
force-pushed
the
subrange-dialect-ops
branch
from
d123ad6 to
70b35e9
Compare
Rajveer100
force-pushed
the
subrange-dialect-ops
branch
2 times, most recently
from
fe978d4 to
5ffba9a
Compare
…ct class. Part of llvm#111591 Currently we have `MLIRContext::getRegisteredOperations` which returns all operations for the given context, with the addition of `MLIRContext::getRegisteredOperationsByDialect` we can now retrieve the same for a given dialect class.
Rajveer100
force-pushed
the
subrange-dialect-ops
branch
from
5ffba9a to
db49a12
Compare
Member
Author
|
Thanks for the approval, could you land this for me?! |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
5 participants
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Part of #111591
Currently we have
MLIRContext::getRegisteredOperationswhich returns all operations for the given context, with the addition ofMLIRContext::getRegisteredOperationsByDialectwe can now retrieve the same for a given dialect class.