[Mlir-commits] [flang] [mlir] [mlir] Allow dialect registry updates in a single-threaded context (PR #101119)
Nikhil Kalra
llvmlistbot at llvm.org
Wed Jul 31 19:37:42 PDT 2024
https://github.com/nikalra updated https://github.com/llvm/llvm-project/pull/101119
>From c6532793c01b129ff7cbf6cf0db86d8ca7325069 Mon Sep 17 00:00:00 2001
From: Nikhil Kalra <nkalra at apple.com>
Date: Wed, 31 Jul 2024 19:24:14 -0700
Subject: [PATCH 1/3] DialectRegistry: add extension ID
---
mlir/include/mlir/IR/DialectRegistry.h | 32 ++++++++++++++++++--------
mlir/lib/IR/Dialect.cpp | 4 ++--
2 files changed, 24 insertions(+), 12 deletions(-)
diff --git a/mlir/include/mlir/IR/DialectRegistry.h b/mlir/include/mlir/IR/DialectRegistry.h
index c13a1a1858eb1..ce5a766843228 100644
--- a/mlir/include/mlir/IR/DialectRegistry.h
+++ b/mlir/include/mlir/IR/DialectRegistry.h
@@ -187,7 +187,8 @@ class DialectRegistry {
nameAndRegistrationIt.second.second);
// Merge the extensions.
for (const auto &extension : extensions)
- destination.extensions.push_back(extension->clone());
+ destination.extensions.emplace_back(extension.first,
+ extension.second->clone());
}
/// Return the names of dialects known to this registry.
@@ -206,39 +207,47 @@ class DialectRegistry {
void applyExtensions(MLIRContext *ctx) const;
/// Add the given extension to the registry.
- void addExtension(std::unique_ptr<DialectExtensionBase> extension) {
- extensions.push_back(std::move(extension));
+ void addExtension(StringRef extensionID,
+ std::unique_ptr<DialectExtensionBase> extension) {
+ extensions.emplace_back(extensionID, std::move(extension));
}
/// Add the given extensions to the registry.
template <typename... ExtensionsT>
void addExtensions() {
- (addExtension(std::make_unique<ExtensionsT>()), ...);
+ (addExtension(ExtensionsT::extensionID, std::make_unique<ExtensionsT>()),
+ ...);
}
/// Add an extension function that requires the given dialects.
/// Note: This bare functor overload is provided in addition to the
/// std::function variant to enable dialect type deduction, e.g.:
- /// registry.addExtension(+[](MLIRContext *ctx, MyDialect *dialect) { ... })
+ /// registry.addExtension("ID", +[](MLIRContext *ctx, MyDialect *dialect) {
+ /// ... })
///
/// is equivalent to:
/// registry.addExtension<MyDialect>(
+ /// "ID",
/// [](MLIRContext *ctx, MyDialect *dialect){ ... }
/// )
template <typename... DialectsT>
- void addExtension(void (*extensionFn)(MLIRContext *, DialectsT *...)) {
+ void addExtension(StringRef extensionID,
+ void (*extensionFn)(MLIRContext *, DialectsT *...)) {
addExtension<DialectsT...>(
+ extensionID,
std::function<void(MLIRContext *, DialectsT * ...)>(extensionFn));
}
template <typename... DialectsT>
void
- addExtension(std::function<void(MLIRContext *, DialectsT *...)> extensionFn) {
+ addExtension(StringRef extensionID,
+ std::function<void(MLIRContext *, DialectsT *...)> extensionFn) {
using ExtensionFnT = std::function<void(MLIRContext *, DialectsT * ...)>;
struct Extension : public DialectExtension<Extension, DialectsT...> {
Extension(const Extension &) = default;
Extension(ExtensionFnT extensionFn)
- : extensionFn(std::move(extensionFn)) {}
+ : DialectExtension<Extension, DialectsT...>(),
+ extensionFn(std::move(extensionFn)) {}
~Extension() override = default;
void apply(MLIRContext *context, DialectsT *...dialects) const final {
@@ -246,7 +255,8 @@ class DialectRegistry {
}
ExtensionFnT extensionFn;
};
- addExtension(std::make_unique<Extension>(std::move(extensionFn)));
+ addExtension(extensionID,
+ std::make_unique<Extension>(std::move(extensionFn)));
}
/// Returns true if the current registry is a subset of 'rhs', i.e. if 'rhs'
@@ -255,7 +265,9 @@ class DialectRegistry {
private:
MapTy registry;
- std::vector<std::unique_ptr<DialectExtensionBase>> extensions;
+ using KeyExtensionPair =
+ std::pair<llvm::StringRef, std::unique_ptr<DialectExtensionBase>>;
+ llvm::SmallVector<KeyExtensionPair> extensions;
};
} // namespace mlir
diff --git a/mlir/lib/IR/Dialect.cpp b/mlir/lib/IR/Dialect.cpp
index 965386681f270..0bdc34bbc5ae8 100644
--- a/mlir/lib/IR/Dialect.cpp
+++ b/mlir/lib/IR/Dialect.cpp
@@ -260,7 +260,7 @@ void DialectRegistry::applyExtensions(Dialect *dialect) const {
// Note: Additional extensions may be added while applying an extension.
for (int i = 0; i < static_cast<int>(extensions.size()); ++i)
- applyExtension(*extensions[i]);
+ applyExtension(*extensions[i].second);
}
void DialectRegistry::applyExtensions(MLIRContext *ctx) const {
@@ -287,7 +287,7 @@ void DialectRegistry::applyExtensions(MLIRContext *ctx) const {
// Note: Additional extensions may be added while applying an extension.
for (int i = 0; i < static_cast<int>(extensions.size()); ++i)
- applyExtension(*extensions[i]);
+ applyExtension(*extensions[i].second);
}
bool DialectRegistry::isSubsetOf(const DialectRegistry &rhs) const {
>From 26669bd8a53297dfe545b407d438c85d0ce32bfd Mon Sep 17 00:00:00 2001
From: Nikhil Kalra <nkalra at apple.com>
Date: Wed, 31 Jul 2024 19:36:25 -0700
Subject: [PATCH 2/3] attach extension IDs
---
flang/lib/Optimizer/Dialect/FIRDialect.cpp | 3 +-
.../transform/Ch2/lib/MyExtension.cpp | 3 +
.../transform/Ch3/lib/MyExtension.cpp | 3 +
.../transform/Ch4/lib/MyExtension.cpp | 3 +
.../Conversion/ArithToLLVM/ArithToLLVM.cpp | 7 +-
.../ComplexToLLVM/ComplexToLLVM.cpp | 1 +
.../ControlFlowToLLVM/ControlFlowToLLVM.cpp | 7 +-
.../ConvertToLLVM/ConvertToLLVMPass.cpp | 5 ++
mlir/lib/Conversion/FuncToLLVM/FuncToLLVM.cpp | 7 +-
.../Conversion/IndexToLLVM/IndexToLLVM.cpp | 7 +-
mlir/lib/Conversion/MathToLLVM/MathToLLVM.cpp | 7 +-
.../Conversion/MemRefToLLVM/MemRefToLLVM.cpp | 7 +-
mlir/lib/Conversion/NVVMToLLVM/NVVMToLLVM.cpp | 7 +-
mlir/lib/Conversion/UBToLLVM/UBToLLVM.cpp | 7 +-
.../Affine/IR/ValueBoundsOpInterfaceImpl.cpp | 11 +--
.../TransformOps/AffineTransformOps.cpp | 2 +
.../Arith/IR/ValueBoundsOpInterfaceImpl.cpp | 16 +++--
.../BufferDeallocationOpInterfaceImpl.cpp | 7 +-
.../BufferViewFlowOpInterfaceImpl.cpp | 7 +-
.../BufferizableOpInterfaceImpl.cpp | 11 +--
.../BufferizationTransformOps.cpp | 2 +
.../FuncBufferizableOpInterfaceImpl.cpp | 14 ++--
.../BufferDeallocationOpInterfaceImpl.cpp | 7 +-
.../BufferizableOpInterfaceImpl.cpp | 10 +--
.../Func/Extensions/InlinerExtension.cpp | 11 +--
.../Extensions/MeshShardingExtensions.cpp | 11 +--
.../Func/TransformOps/FuncTransformOps.cpp | 2 +
.../GPU/TransformOps/GPUTransformOps.cpp | 2 +
.../BufferDeallocationOpInterfaceImpl.cpp | 7 +-
.../Linalg/IR/ValueBoundsOpInterfaceImpl.cpp | 13 ++--
.../Linalg/TransformOps/DialectExtension.cpp | 2 +
.../BufferizableOpInterfaceImpl.cpp | 19 ++---
.../Transforms/MeshShardingInterfaceImpl.cpp | 25 +++----
.../Transforms/RuntimeOpVerification.cpp | 17 ++---
.../SubsetInsertionOpInterfaceImpl.cpp | 9 +--
.../Linalg/Transforms/TilingInterfaceImpl.cpp | 11 +--
.../BufferizableOpInterfaceImpl.cpp | 11 +--
.../Dialect/MemRef/IR/MemRefMemorySlot.cpp | 9 ++-
.../MemRef/IR/ValueBoundsOpInterfaceImpl.cpp | 25 ++++---
.../TransformOps/MemRefTransformOps.cpp | 2 +
.../Transforms/AllocationOpInterfaceImpl.cpp | 13 ++--
.../BufferViewFlowOpInterfaceImpl.cpp | 8 ++-
.../Transforms/RuntimeOpVerification.cpp | 26 +++----
.../NVGPU/TransformOps/NVGPUTransformOps.cpp | 2 +
.../SCF/IR/ValueBoundsOpInterfaceImpl.cpp | 9 +--
.../SCF/TransformOps/SCFTransformOps.cpp | 2 +
.../BufferDeallocationOpInterfaceImpl.cpp | 9 +--
.../BufferizableOpInterfaceImpl.cpp | 23 +++---
.../BufferizableOpInterfaceImpl.cpp | 10 +--
.../TransformOps/SparseTensorTransformOps.cpp | 2 +
.../BufferizableOpInterfaceImpl.cpp | 40 ++++++-----
.../IR/TensorInferTypeOpInterfaceImpl.cpp | 15 ++--
.../Tensor/IR/TensorTilingInterfaceImpl.cpp | 21 +++---
.../Tensor/IR/ValueBoundsOpInterfaceImpl.cpp | 25 ++++---
.../TransformOps/TensorTransformOps.cpp | 23 +++---
.../BufferizableOpInterfaceImpl.cpp | 46 ++++++------
.../SubsetInsertionOpInterfaceImpl.cpp | 33 ++++-----
.../Dialect/Tosa/IR/ShardingInterfaceImpl.cpp | 22 +++---
.../DebugExtension/DebugExtension.cpp | 2 +
.../Transform/IRDLExtension/IRDLExtension.cpp | 2 +
.../Transform/LoopExtension/LoopExtension.cpp | 2 +
.../Transform/PDLExtension/PDLExtension.cpp | 2 +
.../Vector/IR/ValueBoundsOpInterfaceImpl.cpp | 10 +--
.../TransformOps/VectorTransformOps.cpp | 2 +
.../BufferizableOpInterfaceImpl.cpp | 16 +++--
.../Transforms/SubsetOpInterfaceImpl.cpp | 21 +++---
mlir/lib/Interfaces/CastInterfaces.cpp | 9 +--
mlir/lib/Target/LLVM/NVVM/Target.cpp | 7 +-
mlir/lib/Target/LLVM/ROCDL/Target.cpp | 7 +-
.../Dialect/AMX/AMXToLLVMIRTranslation.cpp | 7 +-
.../ArmNeon/ArmNeonToLLVMIRTranslation.cpp | 1 +
.../ArmSME/ArmSMEToLLVMIRTranslation.cpp | 8 ++-
.../ArmSVE/ArmSVEToLLVMIRTranslation.cpp | 8 ++-
.../Builtin/BuiltinToLLVMIRTranslation.cpp | 7 +-
.../Dialect/GPU/GPUToLLVMIRTranslation.cpp | 7 +-
.../LLVMIR/Dialect/GPU/SelectObjectAttr.cpp | 7 +-
.../LLVMIR/LLVMIRToLLVMTranslation.cpp | 7 +-
.../LLVMIR/LLVMToLLVMIRTranslation.cpp | 7 +-
.../Dialect/NVVM/LLVMIRToNVVMTranslation.cpp | 7 +-
.../Dialect/NVVM/NVVMToLLVMIRTranslation.cpp | 7 +-
.../OpenACC/OpenACCToLLVMIRTranslation.cpp | 7 +-
.../OpenMP/OpenMPToLLVMIRTranslation.cpp | 7 +-
.../ROCDL/ROCDLToLLVMIRTranslation.cpp | 7 +-
.../Dialect/VCIX/VCIXToLLVMIRTranslation.cpp | 7 +-
.../X86VectorToLLVMIRTranslation.cpp | 1 +
mlir/lib/Target/SPIRV/Target.cpp | 7 +-
.../Test/TestFromLLVMIRTranslation.cpp | 1 +
.../Dialect/Test/TestToLLVMIRTranslation.cpp | 1 +
.../TestTransformDialectExtension.cpp | 2 +
.../TestTilingInterfaceTransformOps.cpp | 2 +
.../Transform/BuildOnlyExtensionTest.cpp | 2 +
mlir/unittests/IR/DialectTest.cpp | 25 +++++--
mlir/unittests/IR/InterfaceAttachmentTest.cpp | 72 +++++++++++--------
93 files changed, 567 insertions(+), 393 deletions(-)
diff --git a/flang/lib/Optimizer/Dialect/FIRDialect.cpp b/flang/lib/Optimizer/Dialect/FIRDialect.cpp
index 4b1dadaac6728..3bf8b43c3572e 100644
--- a/flang/lib/Optimizer/Dialect/FIRDialect.cpp
+++ b/flang/lib/Optimizer/Dialect/FIRDialect.cpp
@@ -74,7 +74,7 @@ void fir::FIROpsDialect::initialize() {
// Register the FIRInlinerInterface to FIROpsDialect
void fir::addFIRInlinerExtension(mlir::DialectRegistry ®istry) {
registry.addExtension(
- +[](mlir::MLIRContext *ctx, fir::FIROpsDialect *dialect) {
+ "FIR_INLINER", +[](mlir::MLIRContext *ctx, fir::FIROpsDialect *dialect) {
dialect->addInterface<FIRInlinerInterface>();
});
}
@@ -90,6 +90,7 @@ void fir::addFIRInlinerExtension(mlir::DialectRegistry ®istry) {
// when more sophisticated translation is required.
void fir::addFIRToLLVMIRExtension(mlir::DialectRegistry ®istry) {
registry.addExtension(
+ "FIR_LLVM_TRANSLATION",
+[](mlir::MLIRContext *ctx, fir::FIROpsDialect *dialect) {
dialect->addInterface<mlir::LLVMTranslationDialectInterface>();
});
diff --git a/mlir/examples/transform/Ch2/lib/MyExtension.cpp b/mlir/examples/transform/Ch2/lib/MyExtension.cpp
index 68d538a098018..a3f4f9aeb8106 100644
--- a/mlir/examples/transform/Ch2/lib/MyExtension.cpp
+++ b/mlir/examples/transform/Ch2/lib/MyExtension.cpp
@@ -36,6 +36,9 @@ class MyExtension
// dialect definitions. List individual operations and dependent dialects
// here.
void init();
+
+ // Declare a unique ID for this extension.
+ static constexpr llvm::StringRef extensionID = "CH2_TRANSFORM";
};
void MyExtension::init() {
diff --git a/mlir/examples/transform/Ch3/lib/MyExtension.cpp b/mlir/examples/transform/Ch3/lib/MyExtension.cpp
index f7a99423a52ee..0eeeb3b4071d3 100644
--- a/mlir/examples/transform/Ch3/lib/MyExtension.cpp
+++ b/mlir/examples/transform/Ch3/lib/MyExtension.cpp
@@ -42,6 +42,9 @@ class MyExtension
// dialect definitions. List individual operations and dependent dialects
// here.
void init();
+
+ // Declare a unique ID for this extension.
+ static constexpr llvm::StringRef extensionID = "CH3_TRANSFORM";
};
void MyExtension::init() {
diff --git a/mlir/examples/transform/Ch4/lib/MyExtension.cpp b/mlir/examples/transform/Ch4/lib/MyExtension.cpp
index 38c8ca1125a24..1d440f588e60b 100644
--- a/mlir/examples/transform/Ch4/lib/MyExtension.cpp
+++ b/mlir/examples/transform/Ch4/lib/MyExtension.cpp
@@ -38,6 +38,9 @@ class MyExtension
// dialect definitions. List individual operations and dependent dialects
// here.
void init();
+
+ // Declare a unique ID for this extension.
+ static constexpr llvm::StringRef extensionID = "CH4_TRANSFORM";
};
void MyExtension::init() {
diff --git a/mlir/lib/Conversion/ArithToLLVM/ArithToLLVM.cpp b/mlir/lib/Conversion/ArithToLLVM/ArithToLLVM.cpp
index d882f1184f457..d2794a24cd05a 100644
--- a/mlir/lib/Conversion/ArithToLLVM/ArithToLLVM.cpp
+++ b/mlir/lib/Conversion/ArithToLLVM/ArithToLLVM.cpp
@@ -510,9 +510,10 @@ struct ArithToLLVMDialectInterface : public ConvertToLLVMPatternInterface {
void mlir::arith::registerConvertArithToLLVMInterface(
DialectRegistry ®istry) {
- registry.addExtension(+[](MLIRContext *ctx, arith::ArithDialect *dialect) {
- dialect->addInterfaces<ArithToLLVMDialectInterface>();
- });
+ registry.addExtension(
+ "ARITH_TO_LLVM", +[](MLIRContext *ctx, arith::ArithDialect *dialect) {
+ dialect->addInterfaces<ArithToLLVMDialectInterface>();
+ });
}
//===----------------------------------------------------------------------===//
diff --git a/mlir/lib/Conversion/ComplexToLLVM/ComplexToLLVM.cpp b/mlir/lib/Conversion/ComplexToLLVM/ComplexToLLVM.cpp
index 0a3c3a330ff69..3701d62efba50 100644
--- a/mlir/lib/Conversion/ComplexToLLVM/ComplexToLLVM.cpp
+++ b/mlir/lib/Conversion/ComplexToLLVM/ComplexToLLVM.cpp
@@ -385,6 +385,7 @@ struct ComplexToLLVMDialectInterface : public ConvertToLLVMPatternInterface {
void mlir::registerConvertComplexToLLVMInterface(DialectRegistry ®istry) {
registry.addExtension(
+ "COMPLEX_TO_LLVM",
+[](MLIRContext *ctx, complex::ComplexDialect *dialect) {
dialect->addInterfaces<ComplexToLLVMDialectInterface>();
});
diff --git a/mlir/lib/Conversion/ControlFlowToLLVM/ControlFlowToLLVM.cpp b/mlir/lib/Conversion/ControlFlowToLLVM/ControlFlowToLLVM.cpp
index b8e5aec25286d..4a154554c30bc 100644
--- a/mlir/lib/Conversion/ControlFlowToLLVM/ControlFlowToLLVM.cpp
+++ b/mlir/lib/Conversion/ControlFlowToLLVM/ControlFlowToLLVM.cpp
@@ -273,7 +273,8 @@ struct ControlFlowToLLVMDialectInterface
void mlir::cf::registerConvertControlFlowToLLVMInterface(
DialectRegistry ®istry) {
- registry.addExtension(+[](MLIRContext *ctx, cf::ControlFlowDialect *dialect) {
- dialect->addInterfaces<ControlFlowToLLVMDialectInterface>();
- });
+ registry.addExtension(
+ "CF_TO_LLVM", +[](MLIRContext *ctx, cf::ControlFlowDialect *dialect) {
+ dialect->addInterfaces<ControlFlowToLLVMDialectInterface>();
+ });
}
diff --git a/mlir/lib/Conversion/ConvertToLLVM/ConvertToLLVMPass.cpp b/mlir/lib/Conversion/ConvertToLLVM/ConvertToLLVMPass.cpp
index 6135117348a5b..c55607f5e7c60 100644
--- a/mlir/lib/Conversion/ConvertToLLVM/ConvertToLLVMPass.cpp
+++ b/mlir/lib/Conversion/ConvertToLLVM/ConvertToLLVMPass.cpp
@@ -15,6 +15,7 @@
#include "mlir/Pass/Pass.h"
#include "mlir/Rewrite/FrozenRewritePatternSet.h"
#include "mlir/Transforms/DialectConversion.h"
+#include "llvm/ADT/StringRef.h"
#include <memory>
#define DEBUG_TYPE "convert-to-llvm"
@@ -54,6 +55,10 @@ class LoadDependentDialectExtension : public DialectExtensionBase {
std::unique_ptr<DialectExtensionBase> clone() const final {
return std::make_unique<LoadDependentDialectExtension>(*this);
}
+
+ /// Unique ID
+ constexpr static llvm::StringRef extensionID =
+ "CONVERT_TO_LLVM_LOAD_DEPENDENT";
};
/// This is a generic pass to convert to LLVM, it uses the
diff --git a/mlir/lib/Conversion/FuncToLLVM/FuncToLLVM.cpp b/mlir/lib/Conversion/FuncToLLVM/FuncToLLVM.cpp
index c1f6d8bc5b361..5994a4c5fdf0f 100644
--- a/mlir/lib/Conversion/FuncToLLVM/FuncToLLVM.cpp
+++ b/mlir/lib/Conversion/FuncToLLVM/FuncToLLVM.cpp
@@ -786,7 +786,8 @@ struct FuncToLLVMDialectInterface : public ConvertToLLVMPatternInterface {
} // namespace
void mlir::registerConvertFuncToLLVMInterface(DialectRegistry ®istry) {
- registry.addExtension(+[](MLIRContext *ctx, func::FuncDialect *dialect) {
- dialect->addInterfaces<FuncToLLVMDialectInterface>();
- });
+ registry.addExtension(
+ "FUNC_TO_LLVM", +[](MLIRContext *ctx, func::FuncDialect *dialect) {
+ dialect->addInterfaces<FuncToLLVMDialectInterface>();
+ });
}
diff --git a/mlir/lib/Conversion/IndexToLLVM/IndexToLLVM.cpp b/mlir/lib/Conversion/IndexToLLVM/IndexToLLVM.cpp
index 9d8a5d8a0e1c0..e045605475b78 100644
--- a/mlir/lib/Conversion/IndexToLLVM/IndexToLLVM.cpp
+++ b/mlir/lib/Conversion/IndexToLLVM/IndexToLLVM.cpp
@@ -392,7 +392,8 @@ struct IndexToLLVMDialectInterface : public ConvertToLLVMPatternInterface {
void mlir::index::registerConvertIndexToLLVMInterface(
DialectRegistry ®istry) {
- registry.addExtension(+[](MLIRContext *ctx, index::IndexDialect *dialect) {
- dialect->addInterfaces<IndexToLLVMDialectInterface>();
- });
+ registry.addExtension(
+ "INDEX_TO_LLVM", +[](MLIRContext *ctx, index::IndexDialect *dialect) {
+ dialect->addInterfaces<IndexToLLVMDialectInterface>();
+ });
}
diff --git a/mlir/lib/Conversion/MathToLLVM/MathToLLVM.cpp b/mlir/lib/Conversion/MathToLLVM/MathToLLVM.cpp
index 23e957288eb95..f43ce2800e310 100644
--- a/mlir/lib/Conversion/MathToLLVM/MathToLLVM.cpp
+++ b/mlir/lib/Conversion/MathToLLVM/MathToLLVM.cpp
@@ -356,7 +356,8 @@ struct MathToLLVMDialectInterface : public ConvertToLLVMPatternInterface {
} // namespace
void mlir::registerConvertMathToLLVMInterface(DialectRegistry ®istry) {
- registry.addExtension(+[](MLIRContext *ctx, math::MathDialect *dialect) {
- dialect->addInterfaces<MathToLLVMDialectInterface>();
- });
+ registry.addExtension(
+ "MATH_TO_LLVM", +[](MLIRContext *ctx, math::MathDialect *dialect) {
+ dialect->addInterfaces<MathToLLVMDialectInterface>();
+ });
}
diff --git a/mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp b/mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp
index 054827d40f0f3..d0d4b159ddd57 100644
--- a/mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp
+++ b/mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp
@@ -1753,7 +1753,8 @@ struct MemRefToLLVMDialectInterface : public ConvertToLLVMPatternInterface {
} // namespace
void mlir::registerConvertMemRefToLLVMInterface(DialectRegistry ®istry) {
- registry.addExtension(+[](MLIRContext *ctx, memref::MemRefDialect *dialect) {
- dialect->addInterfaces<MemRefToLLVMDialectInterface>();
- });
+ registry.addExtension(
+ "MEMREF_TO_LLVM", +[](MLIRContext *ctx, memref::MemRefDialect *dialect) {
+ dialect->addInterfaces<MemRefToLLVMDialectInterface>();
+ });
}
diff --git a/mlir/lib/Conversion/NVVMToLLVM/NVVMToLLVM.cpp b/mlir/lib/Conversion/NVVMToLLVM/NVVMToLLVM.cpp
index 662ee9e483bc5..c6b856b84b10b 100644
--- a/mlir/lib/Conversion/NVVMToLLVM/NVVMToLLVM.cpp
+++ b/mlir/lib/Conversion/NVVMToLLVM/NVVMToLLVM.cpp
@@ -113,7 +113,8 @@ void mlir::populateNVVMToLLVMConversionPatterns(RewritePatternSet &patterns) {
}
void mlir::registerConvertNVVMToLLVMInterface(DialectRegistry ®istry) {
- registry.addExtension(+[](MLIRContext *ctx, NVVMDialect *dialect) {
- dialect->addInterfaces<NVVMToLLVMDialectInterface>();
- });
+ registry.addExtension(
+ "NVVM_TO_LLVM", +[](MLIRContext *ctx, NVVMDialect *dialect) {
+ dialect->addInterfaces<NVVMToLLVMDialectInterface>();
+ });
}
diff --git a/mlir/lib/Conversion/UBToLLVM/UBToLLVM.cpp b/mlir/lib/Conversion/UBToLLVM/UBToLLVM.cpp
index 0051333a35dcd..12150d2e1977b 100644
--- a/mlir/lib/Conversion/UBToLLVM/UBToLLVM.cpp
+++ b/mlir/lib/Conversion/UBToLLVM/UBToLLVM.cpp
@@ -119,7 +119,8 @@ struct UBToLLVMDialectInterface : public ConvertToLLVMPatternInterface {
} // namespace
void mlir::ub::registerConvertUBToLLVMInterface(DialectRegistry ®istry) {
- registry.addExtension(+[](MLIRContext *ctx, ub::UBDialect *dialect) {
- dialect->addInterfaces<UBToLLVMDialectInterface>();
- });
+ registry.addExtension(
+ "UB_TO_LLVM", +[](MLIRContext *ctx, ub::UBDialect *dialect) {
+ dialect->addInterfaces<UBToLLVMDialectInterface>();
+ });
}
diff --git a/mlir/lib/Dialect/Affine/IR/ValueBoundsOpInterfaceImpl.cpp b/mlir/lib/Dialect/Affine/IR/ValueBoundsOpInterfaceImpl.cpp
index 82a9fb0d49088..2fe39c757fcd3 100644
--- a/mlir/lib/Dialect/Affine/IR/ValueBoundsOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Affine/IR/ValueBoundsOpInterfaceImpl.cpp
@@ -96,11 +96,12 @@ struct AffineMaxOpInterface
void mlir::affine::registerValueBoundsOpInterfaceExternalModels(
DialectRegistry ®istry) {
- registry.addExtension(+[](MLIRContext *ctx, AffineDialect *dialect) {
- AffineApplyOp::attachInterface<AffineApplyOpInterface>(*ctx);
- AffineMaxOp::attachInterface<AffineMaxOpInterface>(*ctx);
- AffineMinOp::attachInterface<AffineMinOpInterface>(*ctx);
- });
+ registry.addExtension(
+ "AFFINE_VALUE_BOUNDS", +[](MLIRContext *ctx, AffineDialect *dialect) {
+ AffineApplyOp::attachInterface<AffineApplyOpInterface>(*ctx);
+ AffineMaxOp::attachInterface<AffineMaxOpInterface>(*ctx);
+ AffineMinOp::attachInterface<AffineMinOpInterface>(*ctx);
+ });
}
FailureOr<int64_t>
diff --git a/mlir/lib/Dialect/Affine/TransformOps/AffineTransformOps.cpp b/mlir/lib/Dialect/Affine/TransformOps/AffineTransformOps.cpp
index 6457655cfe416..c6f5cfd47b046 100644
--- a/mlir/lib/Dialect/Affine/TransformOps/AffineTransformOps.cpp
+++ b/mlir/lib/Dialect/Affine/TransformOps/AffineTransformOps.cpp
@@ -167,6 +167,8 @@ class AffineTransformDialectExtension
#include "mlir/Dialect/Affine/TransformOps/AffineTransformOps.cpp.inc"
>();
}
+
+ static constexpr llvm::StringRef extensionID = "AFFINE_TRANSFORM";
};
} // namespace
diff --git a/mlir/lib/Dialect/Arith/IR/ValueBoundsOpInterfaceImpl.cpp b/mlir/lib/Dialect/Arith/IR/ValueBoundsOpInterfaceImpl.cpp
index 7cfcc4180539c..981a52090ad66 100644
--- a/mlir/lib/Dialect/Arith/IR/ValueBoundsOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Arith/IR/ValueBoundsOpInterfaceImpl.cpp
@@ -150,11 +150,13 @@ struct SelectOpInterface
void mlir::arith::registerValueBoundsOpInterfaceExternalModels(
DialectRegistry ®istry) {
- registry.addExtension(+[](MLIRContext *ctx, arith::ArithDialect *dialect) {
- arith::AddIOp::attachInterface<arith::AddIOpInterface>(*ctx);
- arith::ConstantOp::attachInterface<arith::ConstantOpInterface>(*ctx);
- arith::SubIOp::attachInterface<arith::SubIOpInterface>(*ctx);
- arith::MulIOp::attachInterface<arith::MulIOpInterface>(*ctx);
- arith::SelectOp::attachInterface<arith::SelectOpInterface>(*ctx);
- });
+ registry.addExtension(
+ "ARITH_VALUE_BOUNDS",
+ +[](MLIRContext *ctx, arith::ArithDialect *dialect) {
+ arith::AddIOp::attachInterface<arith::AddIOpInterface>(*ctx);
+ arith::ConstantOp::attachInterface<arith::ConstantOpInterface>(*ctx);
+ arith::SubIOp::attachInterface<arith::SubIOpInterface>(*ctx);
+ arith::MulIOp::attachInterface<arith::MulIOpInterface>(*ctx);
+ arith::SelectOp::attachInterface<arith::SelectOpInterface>(*ctx);
+ });
}
diff --git a/mlir/lib/Dialect/Arith/Transforms/BufferDeallocationOpInterfaceImpl.cpp b/mlir/lib/Dialect/Arith/Transforms/BufferDeallocationOpInterfaceImpl.cpp
index f2e7732e8ea4a..c17f6aacd88f8 100644
--- a/mlir/lib/Dialect/Arith/Transforms/BufferDeallocationOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Arith/Transforms/BufferDeallocationOpInterfaceImpl.cpp
@@ -79,7 +79,8 @@ struct SelectOpInterface
void mlir::arith::registerBufferDeallocationOpInterfaceExternalModels(
DialectRegistry ®istry) {
- registry.addExtension(+[](MLIRContext *ctx, ArithDialect *dialect) {
- SelectOp::attachInterface<SelectOpInterface>(*ctx);
- });
+ registry.addExtension(
+ "ARITH_BUFFER_DEALLOC", +[](MLIRContext *ctx, ArithDialect *dialect) {
+ SelectOp::attachInterface<SelectOpInterface>(*ctx);
+ });
}
diff --git a/mlir/lib/Dialect/Arith/Transforms/BufferViewFlowOpInterfaceImpl.cpp b/mlir/lib/Dialect/Arith/Transforms/BufferViewFlowOpInterfaceImpl.cpp
index 9df9df86b64fb..9f03fa947085f 100644
--- a/mlir/lib/Dialect/Arith/Transforms/BufferViewFlowOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Arith/Transforms/BufferViewFlowOpInterfaceImpl.cpp
@@ -38,7 +38,8 @@ struct SelectOpInterface
void arith::registerBufferViewFlowOpInterfaceExternalModels(
DialectRegistry ®istry) {
- registry.addExtension(+[](MLIRContext *ctx, arith::ArithDialect *dialect) {
- SelectOp::attachInterface<SelectOpInterface>(*ctx);
- });
+ registry.addExtension(
+ "ARITH_BUFFER_FLOW", +[](MLIRContext *ctx, arith::ArithDialect *dialect) {
+ SelectOp::attachInterface<SelectOpInterface>(*ctx);
+ });
}
diff --git a/mlir/lib/Dialect/Arith/Transforms/BufferizableOpInterfaceImpl.cpp b/mlir/lib/Dialect/Arith/Transforms/BufferizableOpInterfaceImpl.cpp
index 5e69a98db8f1e..74b55c124e480 100644
--- a/mlir/lib/Dialect/Arith/Transforms/BufferizableOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Arith/Transforms/BufferizableOpInterfaceImpl.cpp
@@ -206,9 +206,10 @@ struct SelectOpInterface
void mlir::arith::registerBufferizableOpInterfaceExternalModels(
DialectRegistry ®istry) {
- registry.addExtension(+[](MLIRContext *ctx, ArithDialect *dialect) {
- ConstantOp::attachInterface<ConstantOpInterface>(*ctx);
- IndexCastOp::attachInterface<IndexCastOpInterface>(*ctx);
- SelectOp::attachInterface<SelectOpInterface>(*ctx);
- });
+ registry.addExtension(
+ "ARITH_BUFFERIZATION", +[](MLIRContext *ctx, ArithDialect *dialect) {
+ ConstantOp::attachInterface<ConstantOpInterface>(*ctx);
+ IndexCastOp::attachInterface<IndexCastOpInterface>(*ctx);
+ SelectOp::attachInterface<SelectOpInterface>(*ctx);
+ });
}
diff --git a/mlir/lib/Dialect/Bufferization/TransformOps/BufferizationTransformOps.cpp b/mlir/lib/Dialect/Bufferization/TransformOps/BufferizationTransformOps.cpp
index e10c7bd914e35..d0993a5b85e8d 100644
--- a/mlir/lib/Dialect/Bufferization/TransformOps/BufferizationTransformOps.cpp
+++ b/mlir/lib/Dialect/Bufferization/TransformOps/BufferizationTransformOps.cpp
@@ -161,6 +161,8 @@ class BufferizationTransformDialectExtension
#include "mlir/Dialect/Bufferization/TransformOps/BufferizationTransformOps.cpp.inc"
>();
}
+
+ static constexpr llvm::StringRef extensionID = "BUFFERIZATION_TRANSFORM";
};
} // namespace
diff --git a/mlir/lib/Dialect/Bufferization/Transforms/FuncBufferizableOpInterfaceImpl.cpp b/mlir/lib/Dialect/Bufferization/Transforms/FuncBufferizableOpInterfaceImpl.cpp
index 053ea7935260a..6f2ec90bc866e 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/FuncBufferizableOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/FuncBufferizableOpInterfaceImpl.cpp
@@ -82,7 +82,8 @@ getBufferizedFunctionArgType(FuncOp funcOp, int64_t index,
/// Return the FuncOp called by `callOp`.
static FuncOp getCalledFunction(CallOpInterface callOp) {
- SymbolRefAttr sym = llvm::dyn_cast_if_present<SymbolRefAttr>(callOp.getCallableForCallee());
+ SymbolRefAttr sym =
+ llvm::dyn_cast_if_present<SymbolRefAttr>(callOp.getCallableForCallee());
if (!sym)
return nullptr;
return dyn_cast_or_null<FuncOp>(
@@ -489,9 +490,10 @@ struct FuncOpInterface
void mlir::bufferization::func_ext::
registerBufferizableOpInterfaceExternalModels(DialectRegistry ®istry) {
- registry.addExtension(+[](MLIRContext *ctx, func::FuncDialect *dialect) {
- func::CallOp::attachInterface<func_ext::CallOpInterface>(*ctx);
- func::FuncOp::attachInterface<func_ext::FuncOpInterface>(*ctx);
- func::ReturnOp::attachInterface<func_ext::ReturnOpInterface>(*ctx);
- });
+ registry.addExtension(
+ "FUNC_BUFFERIZATION", +[](MLIRContext *ctx, func::FuncDialect *dialect) {
+ func::CallOp::attachInterface<func_ext::CallOpInterface>(*ctx);
+ func::FuncOp::attachInterface<func_ext::FuncOpInterface>(*ctx);
+ func::ReturnOp::attachInterface<func_ext::ReturnOpInterface>(*ctx);
+ });
}
diff --git a/mlir/lib/Dialect/ControlFlow/Transforms/BufferDeallocationOpInterfaceImpl.cpp b/mlir/lib/Dialect/ControlFlow/Transforms/BufferDeallocationOpInterfaceImpl.cpp
index 89546da428fa2..f58f21b35547d 100644
--- a/mlir/lib/Dialect/ControlFlow/Transforms/BufferDeallocationOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/ControlFlow/Transforms/BufferDeallocationOpInterfaceImpl.cpp
@@ -157,7 +157,8 @@ struct CondBranchOpInterface
void mlir::cf::registerBufferDeallocationOpInterfaceExternalModels(
DialectRegistry ®istry) {
- registry.addExtension(+[](MLIRContext *ctx, ControlFlowDialect *dialect) {
- CondBranchOp::attachInterface<CondBranchOpInterface>(*ctx);
- });
+ registry.addExtension(
+ "CF_BUFFER_DEALLOC", +[](MLIRContext *ctx, ControlFlowDialect *dialect) {
+ CondBranchOp::attachInterface<CondBranchOpInterface>(*ctx);
+ });
}
diff --git a/mlir/lib/Dialect/ControlFlow/Transforms/BufferizableOpInterfaceImpl.cpp b/mlir/lib/Dialect/ControlFlow/Transforms/BufferizableOpInterfaceImpl.cpp
index 72f4a1a4f4c66..0db95f92dd95b 100644
--- a/mlir/lib/Dialect/ControlFlow/Transforms/BufferizableOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/ControlFlow/Transforms/BufferizableOpInterfaceImpl.cpp
@@ -63,8 +63,10 @@ struct CondBranchOpInterface
void mlir::cf::registerBufferizableOpInterfaceExternalModels(
DialectRegistry ®istry) {
- registry.addExtension(+[](MLIRContext *ctx, cf::ControlFlowDialect *dialect) {
- cf::BranchOp::attachInterface<BranchOpInterface>(*ctx);
- cf::CondBranchOp::attachInterface<CondBranchOpInterface>(*ctx);
- });
+ registry.addExtension(
+ "CF_BUFFERIZATION",
+ +[](MLIRContext *ctx, cf::ControlFlowDialect *dialect) {
+ cf::BranchOp::attachInterface<BranchOpInterface>(*ctx);
+ cf::CondBranchOp::attachInterface<CondBranchOpInterface>(*ctx);
+ });
}
diff --git a/mlir/lib/Dialect/Func/Extensions/InlinerExtension.cpp b/mlir/lib/Dialect/Func/Extensions/InlinerExtension.cpp
index 719a74a29a622..43e998fbc04f8 100644
--- a/mlir/lib/Dialect/Func/Extensions/InlinerExtension.cpp
+++ b/mlir/lib/Dialect/Func/Extensions/InlinerExtension.cpp
@@ -80,10 +80,11 @@ struct FuncInlinerInterface : public DialectInlinerInterface {
//===----------------------------------------------------------------------===//
void mlir::func::registerInlinerExtension(DialectRegistry ®istry) {
- registry.addExtension(+[](MLIRContext *ctx, func::FuncDialect *dialect) {
- dialect->addInterfaces<FuncInlinerInterface>();
+ registry.addExtension(
+ "FUNC_INLINER", +[](MLIRContext *ctx, func::FuncDialect *dialect) {
+ dialect->addInterfaces<FuncInlinerInterface>();
- // The inliner extension relies on the ControlFlow dialect.
- ctx->getOrLoadDialect<cf::ControlFlowDialect>();
- });
+ // The inliner extension relies on the ControlFlow dialect.
+ ctx->getOrLoadDialect<cf::ControlFlowDialect>();
+ });
}
diff --git a/mlir/lib/Dialect/Func/Extensions/MeshShardingExtensions.cpp b/mlir/lib/Dialect/Func/Extensions/MeshShardingExtensions.cpp
index da508cc95bfe1..833362041848b 100644
--- a/mlir/lib/Dialect/Func/Extensions/MeshShardingExtensions.cpp
+++ b/mlir/lib/Dialect/Func/Extensions/MeshShardingExtensions.cpp
@@ -14,11 +14,12 @@
namespace mlir::func {
void registerShardingInterfaceExternalModels(DialectRegistry ®istry) {
- registry.addExtension(+[](MLIRContext *ctx, FuncDialect *dialect) {
- ReturnOp::attachInterface<
- mesh::IndependentParallelIteratorDomainShardingInterface<ReturnOp>>(
- *ctx);
- });
+ registry.addExtension(
+ "FUNC_SHARDING", +[](MLIRContext *ctx, FuncDialect *dialect) {
+ ReturnOp::attachInterface<
+ mesh::IndependentParallelIteratorDomainShardingInterface<ReturnOp>>(
+ *ctx);
+ });
}
} // namespace mlir::func
diff --git a/mlir/lib/Dialect/Func/TransformOps/FuncTransformOps.cpp b/mlir/lib/Dialect/Func/TransformOps/FuncTransformOps.cpp
index b632b25d0cc67..711a2d9685f6d 100644
--- a/mlir/lib/Dialect/Func/TransformOps/FuncTransformOps.cpp
+++ b/mlir/lib/Dialect/Func/TransformOps/FuncTransformOps.cpp
@@ -246,6 +246,8 @@ class FuncTransformDialectExtension
#include "mlir/Dialect/Func/TransformOps/FuncTransformOps.cpp.inc"
>();
}
+
+ static constexpr llvm::StringRef extensionID = "FUNC_TRANSFORM";
};
} // namespace
diff --git a/mlir/lib/Dialect/GPU/TransformOps/GPUTransformOps.cpp b/mlir/lib/Dialect/GPU/TransformOps/GPUTransformOps.cpp
index 3661c5dea4525..d6a044698746f 100644
--- a/mlir/lib/Dialect/GPU/TransformOps/GPUTransformOps.cpp
+++ b/mlir/lib/Dialect/GPU/TransformOps/GPUTransformOps.cpp
@@ -933,6 +933,8 @@ class GPUTransformDialectExtension
#include "mlir/Dialect/GPU/TransformOps/GPUTransformOps.cpp.inc"
>();
}
+
+ static constexpr llvm::StringRef extensionID = "GPU_TRANSFORM";
};
} // namespace
diff --git a/mlir/lib/Dialect/GPU/Transforms/BufferDeallocationOpInterfaceImpl.cpp b/mlir/lib/Dialect/GPU/Transforms/BufferDeallocationOpInterfaceImpl.cpp
index 6ccc0a26426c1..ddad79f173628 100644
--- a/mlir/lib/Dialect/GPU/Transforms/BufferDeallocationOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/GPU/Transforms/BufferDeallocationOpInterfaceImpl.cpp
@@ -31,7 +31,8 @@ struct GPUTerminatorOpInterface
void mlir::gpu::registerBufferDeallocationOpInterfaceExternalModels(
DialectRegistry ®istry) {
- registry.addExtension(+[](MLIRContext *ctx, GPUDialect *dialect) {
- gpu::TerminatorOp::attachInterface<GPUTerminatorOpInterface>(*ctx);
- });
+ registry.addExtension(
+ "GPU_BUFFER_DEALLOC", +[](MLIRContext *ctx, GPUDialect *dialect) {
+ gpu::TerminatorOp::attachInterface<GPUTerminatorOpInterface>(*ctx);
+ });
}
diff --git a/mlir/lib/Dialect/Linalg/IR/ValueBoundsOpInterfaceImpl.cpp b/mlir/lib/Dialect/Linalg/IR/ValueBoundsOpInterfaceImpl.cpp
index f56ef485069f8..390d6582c487e 100644
--- a/mlir/lib/Dialect/Linalg/IR/ValueBoundsOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Linalg/IR/ValueBoundsOpInterfaceImpl.cpp
@@ -53,9 +53,12 @@ struct IndexOpInterface
void mlir::linalg::registerValueBoundsOpInterfaceExternalModels(
DialectRegistry ®istry) {
- registry.addExtension(+[](MLIRContext *ctx, linalg::LinalgDialect *dialect) {
- IndexOp::attachInterface<IndexOpInterface>(*ctx);
- // Note: ValueBoundsOpInterface implementation is not required for ops that
- // implement `DestinationStyleOpInterface` (for querying shaped OpResults).
- });
+ registry.addExtension(
+ "LINALG_VALUE_BOUNDS",
+ +[](MLIRContext *ctx, linalg::LinalgDialect *dialect) {
+ IndexOp::attachInterface<IndexOpInterface>(*ctx);
+ // Note: ValueBoundsOpInterface implementation is not required for ops
+ // that implement `DestinationStyleOpInterface` (for querying shaped
+ // OpResults).
+ });
}
diff --git a/mlir/lib/Dialect/Linalg/TransformOps/DialectExtension.cpp b/mlir/lib/Dialect/Linalg/TransformOps/DialectExtension.cpp
index f4244ca962232..01c6814bc2082 100644
--- a/mlir/lib/Dialect/Linalg/TransformOps/DialectExtension.cpp
+++ b/mlir/lib/Dialect/Linalg/TransformOps/DialectExtension.cpp
@@ -52,6 +52,8 @@ class LinalgTransformDialectExtension
#include "mlir/Dialect/Linalg/TransformOps/LinalgMatchOps.cpp.inc"
>();
}
+
+ static constexpr llvm::StringRef extensionID = "LINALG_TRANSFORM";
};
} // namespace
diff --git a/mlir/lib/Dialect/Linalg/Transforms/BufferizableOpInterfaceImpl.cpp b/mlir/lib/Dialect/Linalg/Transforms/BufferizableOpInterfaceImpl.cpp
index be158af09d398..dc3bea56da8de 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/BufferizableOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/BufferizableOpInterfaceImpl.cpp
@@ -195,15 +195,18 @@ struct SoftmaxOpInterface
void mlir::linalg::registerBufferizableOpInterfaceExternalModels(
DialectRegistry ®istry) {
- registry.addExtension(+[](MLIRContext *ctx, linalg::LinalgDialect *dialect) {
- // Register all Linalg structured ops. `LinalgOp` is an interface and it is
- // not possible to attach an external interface to an existing interface.
- // Therefore, attach the `BufferizableOpInterface` to all ops one-by-one.
- LinalgOpInterfaceHelper<
+ registry.addExtension(
+ "LINALG_BUFFERIZATION",
+ +[](MLIRContext *ctx, linalg::LinalgDialect *dialect) {
+ // Register all Linalg structured ops. `LinalgOp` is an interface and it
+ // is not possible to attach an external interface to an existing
+ // interface. Therefore, attach the `BufferizableOpInterface` to all ops
+ // one-by-one.
+ LinalgOpInterfaceHelper<
#define GET_OP_LIST
#include "mlir/Dialect/Linalg/IR/LinalgStructuredOps.cpp.inc"
- >::registerOpInterface(ctx);
+ >::registerOpInterface(ctx);
- SoftmaxOp::attachInterface<SoftmaxOpInterface>(*ctx);
- });
+ SoftmaxOp::attachInterface<SoftmaxOpInterface>(*ctx);
+ });
}
diff --git a/mlir/lib/Dialect/Linalg/Transforms/MeshShardingInterfaceImpl.cpp b/mlir/lib/Dialect/Linalg/Transforms/MeshShardingInterfaceImpl.cpp
index 36b6088b83cc2..afb5ffcf31613 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/MeshShardingInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/MeshShardingInterfaceImpl.cpp
@@ -347,20 +347,21 @@ static void registerAll(MLIRContext *ctx) {
}
void registerMeshShardingInterfaceExternalModels(DialectRegistry ®istry) {
- registry.addExtension(+[](MLIRContext *ctx, LinalgDialect *dialect) {
- DialectRegistry registry;
- registry.insert<affine::AffineDialect, arith::ArithDialect, scf::SCFDialect,
- tensor::TensorDialect>();
- ctx->appendDialectRegistry(registry);
- for (StringRef name : registry.getDialectNames())
- ctx->getOrLoadDialect(name);
-
- registerOne<linalg::GenericOp>(ctx);
- registerAll<
+ registry.addExtension(
+ "LINALG_SHARDING", +[](MLIRContext *ctx, LinalgDialect *dialect) {
+ DialectRegistry registry;
+ registry.insert<affine::AffineDialect, arith::ArithDialect,
+ scf::SCFDialect, tensor::TensorDialect>();
+ ctx->appendDialectRegistry(registry);
+ for (StringRef name : registry.getDialectNames())
+ ctx->getOrLoadDialect(name);
+
+ registerOne<linalg::GenericOp>(ctx);
+ registerAll<
#define GET_OP_LIST
#include "mlir/Dialect/Linalg/IR/LinalgStructuredOps.cpp.inc"
- >(ctx);
- });
+ >(ctx);
+ });
}
} // namespace mlir::linalg
diff --git a/mlir/lib/Dialect/Linalg/Transforms/RuntimeOpVerification.cpp b/mlir/lib/Dialect/Linalg/Transforms/RuntimeOpVerification.cpp
index b30182dc84079..6b3ab1a42bd99 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/RuntimeOpVerification.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/RuntimeOpVerification.cpp
@@ -121,15 +121,16 @@ void attachInterface(MLIRContext *ctx) {
void mlir::linalg::registerRuntimeVerifiableOpInterfaceExternalModels(
DialectRegistry ®istry) {
- registry.addExtension(+[](MLIRContext *ctx, LinalgDialect *) {
- attachInterface<
+ registry.addExtension(
+ "LINALG_RUNTIME_VERIFICATION", +[](MLIRContext *ctx, LinalgDialect *) {
+ attachInterface<
#define GET_OP_LIST
#include "mlir/Dialect/Linalg/IR/LinalgStructuredOps.cpp.inc"
- >(ctx);
+ >(ctx);
- // Load additional dialects of which ops may get created.
- ctx->loadDialect<affine::AffineDialect, arith::ArithDialect,
- cf::ControlFlowDialect, index::IndexDialect,
- tensor::TensorDialect>();
- });
+ // Load additional dialects of which ops may get created.
+ ctx->loadDialect<affine::AffineDialect, arith::ArithDialect,
+ cf::ControlFlowDialect, index::IndexDialect,
+ tensor::TensorDialect>();
+ });
}
diff --git a/mlir/lib/Dialect/Linalg/Transforms/SubsetInsertionOpInterfaceImpl.cpp b/mlir/lib/Dialect/Linalg/Transforms/SubsetInsertionOpInterfaceImpl.cpp
index 6fcfa05468eea..04e4b86f8a855 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/SubsetInsertionOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/SubsetInsertionOpInterfaceImpl.cpp
@@ -74,8 +74,9 @@ struct LinalgCopyOpInterface
void mlir::linalg::registerSubsetOpInterfaceExternalModels(
DialectRegistry ®istry) {
- registry.addExtension(+[](MLIRContext *ctx, linalg::LinalgDialect *dialect) {
- linalg::CopyOp::attachInterface<LinalgCopyOpSubsetOpInterface>(*ctx);
- linalg::CopyOp::attachInterface<LinalgCopyOpInterface>(*ctx);
- });
+ registry.addExtension(
+ "LINALG_SUBSET", +[](MLIRContext *ctx, linalg::LinalgDialect *dialect) {
+ linalg::CopyOp::attachInterface<LinalgCopyOpSubsetOpInterface>(*ctx);
+ linalg::CopyOp::attachInterface<LinalgCopyOpInterface>(*ctx);
+ });
}
diff --git a/mlir/lib/Dialect/Linalg/Transforms/TilingInterfaceImpl.cpp b/mlir/lib/Dialect/Linalg/Transforms/TilingInterfaceImpl.cpp
index 2133458efe74c..9cd4b4c6adbd9 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/TilingInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/TilingInterfaceImpl.cpp
@@ -506,10 +506,11 @@ static void registerAll(MLIRContext *ctx) {
void mlir::linalg::registerTilingInterfaceExternalModels(
DialectRegistry ®istry) {
- registry.addExtension(+[](MLIRContext *ctx, linalg::LinalgDialect *dialect) {
- registerOne<linalg::GenericOp>(ctx);
- registerAll<
+ registry.addExtension(
+ "LINALG_TILING", +[](MLIRContext *ctx, linalg::LinalgDialect *dialect) {
+ registerOne<linalg::GenericOp>(ctx);
+ registerAll<
#include "mlir/Dialect/Linalg/IR/LinalgStructuredOps.cpp.inc"
- >(ctx);
- });
+ >(ctx);
+ });
}
diff --git a/mlir/lib/Dialect/MLProgram/Transforms/BufferizableOpInterfaceImpl.cpp b/mlir/lib/Dialect/MLProgram/Transforms/BufferizableOpInterfaceImpl.cpp
index 926d580ac7852..a07a412f11e85 100644
--- a/mlir/lib/Dialect/MLProgram/Transforms/BufferizableOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/MLProgram/Transforms/BufferizableOpInterfaceImpl.cpp
@@ -149,11 +149,12 @@ struct GlobalStoreOpInterface
} // namespace
void registerBufferizableOpInterfaceExternalModels(DialectRegistry ®istry) {
- registry.addExtension(+[](MLIRContext *ctx, MLProgramDialect *) {
- GlobalOp::attachInterface<GlobalOpInterface>(*ctx);
- GlobalLoadOp::attachInterface<GlobalLoadOpInterface>(*ctx);
- GlobalStoreOp::attachInterface<GlobalStoreOpInterface>(*ctx);
- });
+ registry.addExtension(
+ "MLPROGRAM_BUFFERIZATION", +[](MLIRContext *ctx, MLProgramDialect *) {
+ GlobalOp::attachInterface<GlobalOpInterface>(*ctx);
+ GlobalLoadOp::attachInterface<GlobalLoadOpInterface>(*ctx);
+ GlobalStoreOp::attachInterface<GlobalStoreOpInterface>(*ctx);
+ });
}
} // namespace ml_program
} // namespace mlir
diff --git a/mlir/lib/Dialect/MemRef/IR/MemRefMemorySlot.cpp b/mlir/lib/Dialect/MemRef/IR/MemRefMemorySlot.cpp
index f630c48cdcaa1..7c21858adc5fb 100644
--- a/mlir/lib/Dialect/MemRef/IR/MemRefMemorySlot.cpp
+++ b/mlir/lib/Dialect/MemRef/IR/MemRefMemorySlot.cpp
@@ -347,7 +347,10 @@ struct MemRefDestructurableTypeExternalModel
//===----------------------------------------------------------------------===//
void mlir::memref::registerMemorySlotExternalModels(DialectRegistry ®istry) {
- registry.addExtension(+[](MLIRContext *ctx, BuiltinDialect *dialect) {
- MemRefType::attachInterface<MemRefDestructurableTypeExternalModel>(*ctx);
- });
+ registry.addExtension(
+ "MEMREF_MEMORY_SLOT_MODELS",
+ +[](MLIRContext *ctx, BuiltinDialect *dialect) {
+ MemRefType::attachInterface<MemRefDestructurableTypeExternalModel>(
+ *ctx);
+ });
}
diff --git a/mlir/lib/Dialect/MemRef/IR/ValueBoundsOpInterfaceImpl.cpp b/mlir/lib/Dialect/MemRef/IR/ValueBoundsOpInterfaceImpl.cpp
index daec22cf6ebdc..fde96046b3ca2 100644
--- a/mlir/lib/Dialect/MemRef/IR/ValueBoundsOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/MemRef/IR/ValueBoundsOpInterfaceImpl.cpp
@@ -115,15 +115,18 @@ struct SubViewOpInterface
void mlir::memref::registerValueBoundsOpInterfaceExternalModels(
DialectRegistry ®istry) {
- registry.addExtension(+[](MLIRContext *ctx, memref::MemRefDialect *dialect) {
- memref::AllocOp::attachInterface<memref::AllocOpInterface<memref::AllocOp>>(
- *ctx);
- memref::AllocaOp::attachInterface<
- memref::AllocOpInterface<memref::AllocaOp>>(*ctx);
- memref::CastOp::attachInterface<memref::CastOpInterface>(*ctx);
- memref::DimOp::attachInterface<memref::DimOpInterface>(*ctx);
- memref::GetGlobalOp::attachInterface<memref::GetGlobalOpInterface>(*ctx);
- memref::RankOp::attachInterface<memref::RankOpInterface>(*ctx);
- memref::SubViewOp::attachInterface<memref::SubViewOpInterface>(*ctx);
- });
+ registry.addExtension(
+ "MEMREF_VALUE_BOUNDS",
+ +[](MLIRContext *ctx, memref::MemRefDialect *dialect) {
+ memref::AllocOp::attachInterface<
+ memref::AllocOpInterface<memref::AllocOp>>(*ctx);
+ memref::AllocaOp::attachInterface<
+ memref::AllocOpInterface<memref::AllocaOp>>(*ctx);
+ memref::CastOp::attachInterface<memref::CastOpInterface>(*ctx);
+ memref::DimOp::attachInterface<memref::DimOpInterface>(*ctx);
+ memref::GetGlobalOp::attachInterface<memref::GetGlobalOpInterface>(
+ *ctx);
+ memref::RankOp::attachInterface<memref::RankOpInterface>(*ctx);
+ memref::SubViewOp::attachInterface<memref::SubViewOpInterface>(*ctx);
+ });
}
diff --git a/mlir/lib/Dialect/MemRef/TransformOps/MemRefTransformOps.cpp b/mlir/lib/Dialect/MemRef/TransformOps/MemRefTransformOps.cpp
index 8469e84c668cb..20f96a018d330 100644
--- a/mlir/lib/Dialect/MemRef/TransformOps/MemRefTransformOps.cpp
+++ b/mlir/lib/Dialect/MemRef/TransformOps/MemRefTransformOps.cpp
@@ -323,6 +323,8 @@ class MemRefTransformDialectExtension
#include "mlir/Dialect/MemRef/TransformOps/MemRefTransformOps.cpp.inc"
>();
}
+
+ static constexpr llvm::StringRef extensionID = "MEMREF_TRANSFORM";
};
} // namespace
diff --git a/mlir/lib/Dialect/MemRef/Transforms/AllocationOpInterfaceImpl.cpp b/mlir/lib/Dialect/MemRef/Transforms/AllocationOpInterfaceImpl.cpp
index c433415944323..2cde261180e9e 100644
--- a/mlir/lib/Dialect/MemRef/Transforms/AllocationOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/MemRef/Transforms/AllocationOpInterfaceImpl.cpp
@@ -60,10 +60,11 @@ struct DefaultReallocationInterface
void mlir::memref::registerAllocationOpInterfaceExternalModels(
DialectRegistry ®istry) {
- registry.addExtension(+[](MLIRContext *ctx, memref::MemRefDialect *dialect) {
- memref::AllocOp::attachInterface<DefaultAllocationInterface>(*ctx);
- memref::AllocaOp::attachInterface<
- DefaultAutomaticAllocationHoistingInterface>(*ctx);
- memref::ReallocOp::attachInterface<DefaultReallocationInterface>(*ctx);
- });
+ registry.addExtension(
+ "MEMREF_ALLOC", +[](MLIRContext *ctx, memref::MemRefDialect *dialect) {
+ memref::AllocOp::attachInterface<DefaultAllocationInterface>(*ctx);
+ memref::AllocaOp::attachInterface<
+ DefaultAutomaticAllocationHoistingInterface>(*ctx);
+ memref::ReallocOp::attachInterface<DefaultReallocationInterface>(*ctx);
+ });
}
diff --git a/mlir/lib/Dialect/MemRef/Transforms/BufferViewFlowOpInterfaceImpl.cpp b/mlir/lib/Dialect/MemRef/Transforms/BufferViewFlowOpInterfaceImpl.cpp
index bbb269bd00161..5076c0be916c9 100644
--- a/mlir/lib/Dialect/MemRef/Transforms/BufferViewFlowOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/MemRef/Transforms/BufferViewFlowOpInterfaceImpl.cpp
@@ -42,7 +42,9 @@ struct ReallocOpInterface
void memref::registerBufferViewFlowOpInterfaceExternalModels(
DialectRegistry ®istry) {
- registry.addExtension(+[](MLIRContext *ctx, memref::MemRefDialect *dialect) {
- ReallocOp::attachInterface<ReallocOpInterface>(*ctx);
- });
+ registry.addExtension(
+ "MEMREF_BUFFER_FLOW",
+ +[](MLIRContext *ctx, memref::MemRefDialect *dialect) {
+ ReallocOp::attachInterface<ReallocOpInterface>(*ctx);
+ });
}
diff --git a/mlir/lib/Dialect/MemRef/Transforms/RuntimeOpVerification.cpp b/mlir/lib/Dialect/MemRef/Transforms/RuntimeOpVerification.cpp
index 450bfa0cec0c7..80fc16d058b41 100644
--- a/mlir/lib/Dialect/MemRef/Transforms/RuntimeOpVerification.cpp
+++ b/mlir/lib/Dialect/MemRef/Transforms/RuntimeOpVerification.cpp
@@ -333,16 +333,18 @@ struct ExpandShapeOpInterface
void mlir::memref::registerRuntimeVerifiableOpInterfaceExternalModels(
DialectRegistry ®istry) {
- registry.addExtension(+[](MLIRContext *ctx, memref::MemRefDialect *dialect) {
- CastOp::attachInterface<CastOpInterface>(*ctx);
- ExpandShapeOp::attachInterface<ExpandShapeOpInterface>(*ctx);
- LoadOp::attachInterface<LoadStoreOpInterface<LoadOp>>(*ctx);
- ReinterpretCastOp::attachInterface<ReinterpretCastOpInterface>(*ctx);
- StoreOp::attachInterface<LoadStoreOpInterface<StoreOp>>(*ctx);
- SubViewOp::attachInterface<SubViewOpInterface>(*ctx);
-
- // Load additional dialects of which ops may get created.
- ctx->loadDialect<affine::AffineDialect, arith::ArithDialect,
- cf::ControlFlowDialect>();
- });
+ registry.addExtension(
+ "MEMREF_RUNTIME_VERIFICATION",
+ +[](MLIRContext *ctx, memref::MemRefDialect *dialect) {
+ CastOp::attachInterface<CastOpInterface>(*ctx);
+ ExpandShapeOp::attachInterface<ExpandShapeOpInterface>(*ctx);
+ LoadOp::attachInterface<LoadStoreOpInterface<LoadOp>>(*ctx);
+ ReinterpretCastOp::attachInterface<ReinterpretCastOpInterface>(*ctx);
+ StoreOp::attachInterface<LoadStoreOpInterface<StoreOp>>(*ctx);
+ SubViewOp::attachInterface<SubViewOpInterface>(*ctx);
+
+ // Load additional dialects of which ops may get created.
+ ctx->loadDialect<affine::AffineDialect, arith::ArithDialect,
+ cf::ControlFlowDialect>();
+ });
}
diff --git a/mlir/lib/Dialect/NVGPU/TransformOps/NVGPUTransformOps.cpp b/mlir/lib/Dialect/NVGPU/TransformOps/NVGPUTransformOps.cpp
index 7d3d868b326c6..8315ad2f5892b 100644
--- a/mlir/lib/Dialect/NVGPU/TransformOps/NVGPUTransformOps.cpp
+++ b/mlir/lib/Dialect/NVGPU/TransformOps/NVGPUTransformOps.cpp
@@ -1146,6 +1146,8 @@ class NVGPUTransformDialectExtension
#include "mlir/Dialect/NVGPU/TransformOps/NVGPUTransformOps.cpp.inc"
>();
}
+
+ static constexpr llvm::StringRef extensionID = "NVGPU_TRANSFORM";
};
} // namespace
diff --git a/mlir/lib/Dialect/SCF/IR/ValueBoundsOpInterfaceImpl.cpp b/mlir/lib/Dialect/SCF/IR/ValueBoundsOpInterfaceImpl.cpp
index 17a1c016ea16d..c6ede24e605df 100644
--- a/mlir/lib/Dialect/SCF/IR/ValueBoundsOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/SCF/IR/ValueBoundsOpInterfaceImpl.cpp
@@ -159,8 +159,9 @@ struct IfOpInterface
void mlir::scf::registerValueBoundsOpInterfaceExternalModels(
DialectRegistry ®istry) {
- registry.addExtension(+[](MLIRContext *ctx, scf::SCFDialect *dialect) {
- scf::ForOp::attachInterface<scf::ForOpInterface>(*ctx);
- scf::IfOp::attachInterface<scf::IfOpInterface>(*ctx);
- });
+ registry.addExtension(
+ "SCF_VALUE_BOUNDS", +[](MLIRContext *ctx, scf::SCFDialect *dialect) {
+ scf::ForOp::attachInterface<scf::ForOpInterface>(*ctx);
+ scf::IfOp::attachInterface<scf::IfOpInterface>(*ctx);
+ });
}
diff --git a/mlir/lib/Dialect/SCF/TransformOps/SCFTransformOps.cpp b/mlir/lib/Dialect/SCF/TransformOps/SCFTransformOps.cpp
index c4a55c302d0a3..c56d1e46bc0fd 100644
--- a/mlir/lib/Dialect/SCF/TransformOps/SCFTransformOps.cpp
+++ b/mlir/lib/Dialect/SCF/TransformOps/SCFTransformOps.cpp
@@ -624,6 +624,8 @@ class SCFTransformDialectExtension
#include "mlir/Dialect/SCF/TransformOps/SCFTransformOps.cpp.inc"
>();
}
+
+ static constexpr llvm::StringRef extensionID = "SCF_TRANSFORM";
};
} // namespace
diff --git a/mlir/lib/Dialect/SCF/Transforms/BufferDeallocationOpInterfaceImpl.cpp b/mlir/lib/Dialect/SCF/Transforms/BufferDeallocationOpInterfaceImpl.cpp
index 24fbc1dca8361..bea0728e16987 100644
--- a/mlir/lib/Dialect/SCF/Transforms/BufferDeallocationOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/SCF/Transforms/BufferDeallocationOpInterfaceImpl.cpp
@@ -75,8 +75,9 @@ struct ReduceReturnOpInterface
void mlir::scf::registerBufferDeallocationOpInterfaceExternalModels(
DialectRegistry ®istry) {
- registry.addExtension(+[](MLIRContext *ctx, SCFDialect *dialect) {
- InParallelOp::attachInterface<InParallelOpInterface>(*ctx);
- ReduceReturnOp::attachInterface<ReduceReturnOpInterface>(*ctx);
- });
+ registry.addExtension(
+ "SCF_BUFFER_DEALLOC", +[](MLIRContext *ctx, SCFDialect *dialect) {
+ InParallelOp::attachInterface<InParallelOpInterface>(*ctx);
+ ReduceReturnOp::attachInterface<ReduceReturnOpInterface>(*ctx);
+ });
}
diff --git a/mlir/lib/Dialect/SCF/Transforms/BufferizableOpInterfaceImpl.cpp b/mlir/lib/Dialect/SCF/Transforms/BufferizableOpInterfaceImpl.cpp
index cf40443ff3839..0aa9fd4385bb5 100644
--- a/mlir/lib/Dialect/SCF/Transforms/BufferizableOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/SCF/Transforms/BufferizableOpInterfaceImpl.cpp
@@ -1352,15 +1352,16 @@ struct InParallelOpInterface
void mlir::scf::registerBufferizableOpInterfaceExternalModels(
DialectRegistry ®istry) {
- registry.addExtension(+[](MLIRContext *ctx, scf::SCFDialect *dialect) {
- ConditionOp::attachInterface<ConditionOpInterface>(*ctx);
- ExecuteRegionOp::attachInterface<ExecuteRegionOpInterface>(*ctx);
- ForOp::attachInterface<ForOpInterface>(*ctx);
- IfOp::attachInterface<IfOpInterface>(*ctx);
- IndexSwitchOp::attachInterface<IndexSwitchOpInterface>(*ctx);
- ForallOp::attachInterface<ForallOpInterface>(*ctx);
- InParallelOp::attachInterface<InParallelOpInterface>(*ctx);
- WhileOp::attachInterface<WhileOpInterface>(*ctx);
- YieldOp::attachInterface<YieldOpInterface>(*ctx);
- });
+ registry.addExtension(
+ "SCF_BUFFERIZATION", +[](MLIRContext *ctx, scf::SCFDialect *dialect) {
+ ConditionOp::attachInterface<ConditionOpInterface>(*ctx);
+ ExecuteRegionOp::attachInterface<ExecuteRegionOpInterface>(*ctx);
+ ForOp::attachInterface<ForOpInterface>(*ctx);
+ IfOp::attachInterface<IfOpInterface>(*ctx);
+ IndexSwitchOp::attachInterface<IndexSwitchOpInterface>(*ctx);
+ ForallOp::attachInterface<ForallOpInterface>(*ctx);
+ InParallelOp::attachInterface<InParallelOpInterface>(*ctx);
+ WhileOp::attachInterface<WhileOpInterface>(*ctx);
+ YieldOp::attachInterface<YieldOpInterface>(*ctx);
+ });
}
diff --git a/mlir/lib/Dialect/Shape/Transforms/BufferizableOpInterfaceImpl.cpp b/mlir/lib/Dialect/Shape/Transforms/BufferizableOpInterfaceImpl.cpp
index 66a2e45001781..b82c6624d0c32 100644
--- a/mlir/lib/Dialect/Shape/Transforms/BufferizableOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Shape/Transforms/BufferizableOpInterfaceImpl.cpp
@@ -137,8 +137,10 @@ struct AssumingYieldOpInterface
void mlir::shape::registerBufferizableOpInterfaceExternalModels(
DialectRegistry ®istry) {
- registry.addExtension(+[](MLIRContext *ctx, shape::ShapeDialect *dialect) {
- shape::AssumingOp::attachInterface<AssumingOpInterface>(*ctx);
- shape::AssumingYieldOp::attachInterface<AssumingYieldOpInterface>(*ctx);
- });
+ registry.addExtension(
+ "SHAPE_BUFFERIZATION",
+ +[](MLIRContext *ctx, shape::ShapeDialect *dialect) {
+ shape::AssumingOp::attachInterface<AssumingOpInterface>(*ctx);
+ shape::AssumingYieldOp::attachInterface<AssumingYieldOpInterface>(*ctx);
+ });
}
diff --git a/mlir/lib/Dialect/SparseTensor/TransformOps/SparseTensorTransformOps.cpp b/mlir/lib/Dialect/SparseTensor/TransformOps/SparseTensorTransformOps.cpp
index ca19259ebffa6..e2b0c77f5f5a0 100644
--- a/mlir/lib/Dialect/SparseTensor/TransformOps/SparseTensorTransformOps.cpp
+++ b/mlir/lib/Dialect/SparseTensor/TransformOps/SparseTensorTransformOps.cpp
@@ -45,6 +45,8 @@ class SparseTensorTransformDialectExtension
#include "mlir/Dialect/SparseTensor/TransformOps/SparseTensorTransformOps.cpp.inc"
>();
}
+
+ static constexpr llvm::StringRef extensionID = "SPARSE_TENSOR_TRANSFORM";
};
} // namespace
diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/BufferizableOpInterfaceImpl.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/BufferizableOpInterfaceImpl.cpp
index 7734d1d258453..aca1a13a8b913 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/BufferizableOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/BufferizableOpInterfaceImpl.cpp
@@ -326,22 +326,26 @@ struct ToValuesOpInterface
void mlir::sparse_tensor::registerBufferizableOpInterfaceExternalModels(
DialectRegistry ®istry) {
- registry.addExtension(+[](MLIRContext *ctx,
- sparse_tensor::SparseTensorDialect *dialect) {
- sparse_tensor::ConcatenateOp::attachInterface<ConcatenateOpInterface>(*ctx);
- sparse_tensor::ConvertOp::attachInterface<ConvertOpInterface>(*ctx);
- sparse_tensor::LoadOp::attachInterface<LoadOpInterface>(*ctx);
- sparse_tensor::NewOp::attachInterface<NewOpInterface>(*ctx);
- sparse_tensor::NumberOfEntriesOp::attachInterface<
- NumberOfEntriesOpInterface>(*ctx);
- sparse_tensor::AssembleOp::attachInterface<AssembleOpInterface>(*ctx);
- sparse_tensor::DisassembleOp::attachInterface<DisassembleOpInterface>(*ctx);
- sparse_tensor::ForeachOp::attachInterface<ForeachOpInterface>(*ctx);
- sparse_tensor::ToCoordinatesBufferOp::attachInterface<
- ToCoordinatesBufferOpInterface>(*ctx);
- sparse_tensor::ToCoordinatesOp::attachInterface<ToCoordinatesOpInterface>(
- *ctx);
- sparse_tensor::ToPositionsOp::attachInterface<ToPositionsOpInterface>(*ctx);
- sparse_tensor::ToValuesOp::attachInterface<ToValuesOpInterface>(*ctx);
- });
+ registry.addExtension(
+ "SPARSE_TENSOR_BUFFERIZATION",
+ +[](MLIRContext *ctx, sparse_tensor::SparseTensorDialect *dialect) {
+ sparse_tensor::ConcatenateOp::attachInterface<ConcatenateOpInterface>(
+ *ctx);
+ sparse_tensor::ConvertOp::attachInterface<ConvertOpInterface>(*ctx);
+ sparse_tensor::LoadOp::attachInterface<LoadOpInterface>(*ctx);
+ sparse_tensor::NewOp::attachInterface<NewOpInterface>(*ctx);
+ sparse_tensor::NumberOfEntriesOp::attachInterface<
+ NumberOfEntriesOpInterface>(*ctx);
+ sparse_tensor::AssembleOp::attachInterface<AssembleOpInterface>(*ctx);
+ sparse_tensor::DisassembleOp::attachInterface<DisassembleOpInterface>(
+ *ctx);
+ sparse_tensor::ForeachOp::attachInterface<ForeachOpInterface>(*ctx);
+ sparse_tensor::ToCoordinatesBufferOp::attachInterface<
+ ToCoordinatesBufferOpInterface>(*ctx);
+ sparse_tensor::ToCoordinatesOp::attachInterface<
+ ToCoordinatesOpInterface>(*ctx);
+ sparse_tensor::ToPositionsOp::attachInterface<ToPositionsOpInterface>(
+ *ctx);
+ sparse_tensor::ToValuesOp::attachInterface<ToValuesOpInterface>(*ctx);
+ });
}
diff --git a/mlir/lib/Dialect/Tensor/IR/TensorInferTypeOpInterfaceImpl.cpp b/mlir/lib/Dialect/Tensor/IR/TensorInferTypeOpInterfaceImpl.cpp
index 7ff435a033985..48b1dc2bdaa6c 100644
--- a/mlir/lib/Dialect/Tensor/IR/TensorInferTypeOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Tensor/IR/TensorInferTypeOpInterfaceImpl.cpp
@@ -201,11 +201,12 @@ struct ReifyPadOp
void mlir::tensor::registerInferTypeOpInterfaceExternalModels(
DialectRegistry ®istry) {
- registry.addExtension(+[](MLIRContext *ctx, TensorDialect *dialect) {
- ExpandShapeOp::attachInterface<
- ReifyExpandOrCollapseShapeOp<tensor::ExpandShapeOp>>(*ctx);
- CollapseShapeOp::attachInterface<
- ReifyExpandOrCollapseShapeOp<tensor::CollapseShapeOp>>(*ctx);
- PadOp::attachInterface<ReifyPadOp>(*ctx);
- });
+ registry.addExtension(
+ "TENSOR_INFER_TYPE", +[](MLIRContext *ctx, TensorDialect *dialect) {
+ ExpandShapeOp::attachInterface<
+ ReifyExpandOrCollapseShapeOp<tensor::ExpandShapeOp>>(*ctx);
+ CollapseShapeOp::attachInterface<
+ ReifyExpandOrCollapseShapeOp<tensor::CollapseShapeOp>>(*ctx);
+ PadOp::attachInterface<ReifyPadOp>(*ctx);
+ });
}
diff --git a/mlir/lib/Dialect/Tensor/IR/TensorTilingInterfaceImpl.cpp b/mlir/lib/Dialect/Tensor/IR/TensorTilingInterfaceImpl.cpp
index 9b2a97eb2b006..88801be553148 100644
--- a/mlir/lib/Dialect/Tensor/IR/TensorTilingInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Tensor/IR/TensorTilingInterfaceImpl.cpp
@@ -785,17 +785,20 @@ FailureOr<TilingResult> tensor::bubbleUpPadSlice(OpBuilder &b,
void mlir::tensor::registerTilingInterfaceExternalModels(
DialectRegistry ®istry) {
- registry.addExtension(+[](MLIRContext *ctx, TensorDialect *dialect) {
- tensor::PadOp::attachInterface<PadOpTiling>(*ctx);
- tensor::PackOp::attachInterface<PackOpTiling>(*ctx);
- tensor::UnPackOp::attachInterface<UnPackOpTiling>(*ctx);
- });
+ registry.addExtension(
+ "TENSOR_TILING", +[](MLIRContext *ctx, TensorDialect *dialect) {
+ tensor::PadOp::attachInterface<PadOpTiling>(*ctx);
+ tensor::PackOp::attachInterface<PackOpTiling>(*ctx);
+ tensor::UnPackOp::attachInterface<UnPackOpTiling>(*ctx);
+ });
}
void mlir::tensor::registerTilingInterfaceExternalModelsForPackUnPackOps(
DialectRegistry ®istry) {
- registry.addExtension(+[](MLIRContext *ctx, TensorDialect *dialect) {
- tensor::PackOp::attachInterface<PackOpTiling>(*ctx);
- tensor::UnPackOp::attachInterface<UnPackOpTiling>(*ctx);
- });
+ registry.addExtension(
+ "TENOR_TILING_PACK_UNPACK",
+ +[](MLIRContext *ctx, TensorDialect *dialect) {
+ tensor::PackOp::attachInterface<PackOpTiling>(*ctx);
+ tensor::UnPackOp::attachInterface<UnPackOpTiling>(*ctx);
+ });
}
diff --git a/mlir/lib/Dialect/Tensor/IR/ValueBoundsOpInterfaceImpl.cpp b/mlir/lib/Dialect/Tensor/IR/ValueBoundsOpInterfaceImpl.cpp
index 06f2c16406d3c..3b68822a4813d 100644
--- a/mlir/lib/Dialect/Tensor/IR/ValueBoundsOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Tensor/IR/ValueBoundsOpInterfaceImpl.cpp
@@ -114,15 +114,18 @@ struct RankOpInterface
void mlir::tensor::registerValueBoundsOpInterfaceExternalModels(
DialectRegistry ®istry) {
- registry.addExtension(+[](MLIRContext *ctx, tensor::TensorDialect *dialect) {
- tensor::CastOp::attachInterface<tensor::CastOpInterface>(*ctx);
- tensor::DimOp::attachInterface<tensor::DimOpInterface>(*ctx);
- tensor::EmptyOp::attachInterface<tensor::EmptyOpInterface>(*ctx);
- tensor::ExtractSliceOp::attachInterface<tensor::ExtractSliceOpInterface>(
- *ctx);
- tensor::PadOp::attachInterface<tensor::PadOpInterface>(*ctx);
- tensor::RankOp::attachInterface<tensor::RankOpInterface>(*ctx);
- // Note: ValueBoundsOpInterface implementation is not required for ops that
- // implement `DestinationStyleOpInterface` (for querying shaped OpResults).
- });
+ registry.addExtension(
+ "TENSOR_VALUE_BOUNDS",
+ +[](MLIRContext *ctx, tensor::TensorDialect *dialect) {
+ tensor::CastOp::attachInterface<tensor::CastOpInterface>(*ctx);
+ tensor::DimOp::attachInterface<tensor::DimOpInterface>(*ctx);
+ tensor::EmptyOp::attachInterface<tensor::EmptyOpInterface>(*ctx);
+ tensor::ExtractSliceOp::attachInterface<
+ tensor::ExtractSliceOpInterface>(*ctx);
+ tensor::PadOp::attachInterface<tensor::PadOpInterface>(*ctx);
+ tensor::RankOp::attachInterface<tensor::RankOpInterface>(*ctx);
+ // Note: ValueBoundsOpInterface implementation is not required for ops
+ // that implement `DestinationStyleOpInterface` (for querying shaped
+ // OpResults).
+ });
}
diff --git a/mlir/lib/Dialect/Tensor/TransformOps/TensorTransformOps.cpp b/mlir/lib/Dialect/Tensor/TransformOps/TensorTransformOps.cpp
index 33016f84056e9..74f795bf02ba4 100644
--- a/mlir/lib/Dialect/Tensor/TransformOps/TensorTransformOps.cpp
+++ b/mlir/lib/Dialect/Tensor/TransformOps/TensorTransformOps.cpp
@@ -70,15 +70,18 @@ struct ReassociativeReshapeOpReplacementInterface
void tensor::registerFindPayloadReplacementOpInterfaceExternalModels(
DialectRegistry ®istry) {
- registry.addExtension(+[](MLIRContext *ctx, tensor::TensorDialect *dialect) {
- CollapseShapeOp::attachInterface<
- ReassociativeReshapeOpReplacementInterface<CollapseShapeOp>>(*ctx);
- ExpandShapeOp::attachInterface<
- ReassociativeReshapeOpReplacementInterface<ExpandShapeOp>>(*ctx);
- ExtractSliceOp::attachInterface<ExtractSliceOpReplacementInterface>(*ctx);
- InsertSliceOp::attachInterface<InsertSliceOpReplacementInterface>(*ctx);
- ReshapeOp::attachInterface<ReshapeOpReplacementInterface>(*ctx);
- });
+ registry.addExtension(
+ "TENSOR_PAYLOAD_REPLACEMENT",
+ +[](MLIRContext *ctx, tensor::TensorDialect *dialect) {
+ CollapseShapeOp::attachInterface<
+ ReassociativeReshapeOpReplacementInterface<CollapseShapeOp>>(*ctx);
+ ExpandShapeOp::attachInterface<
+ ReassociativeReshapeOpReplacementInterface<ExpandShapeOp>>(*ctx);
+ ExtractSliceOp::attachInterface<ExtractSliceOpReplacementInterface>(
+ *ctx);
+ InsertSliceOp::attachInterface<InsertSliceOpReplacementInterface>(*ctx);
+ ReshapeOp::attachInterface<ReshapeOpReplacementInterface>(*ctx);
+ });
}
//===----------------------------------------------------------------------===//
@@ -247,6 +250,8 @@ class TensorTransformDialectExtension
#include "mlir/Dialect/Tensor/TransformOps/TensorTransformOps.cpp.inc"
>();
}
+
+ static constexpr llvm::StringRef extensionID = "TENSOR_TRANSFORM";
};
} // namespace
diff --git a/mlir/lib/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.cpp b/mlir/lib/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.cpp
index 87464ccb71720..b7f5fe676a7d2 100644
--- a/mlir/lib/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.cpp
@@ -1055,28 +1055,30 @@ struct SplatOpInterface
void mlir::tensor::registerBufferizableOpInterfaceExternalModels(
DialectRegistry ®istry) {
- registry.addExtension(+[](MLIRContext *ctx, tensor::TensorDialect *dialect) {
- CastOp::attachInterface<CastOpInterface>(*ctx);
- CollapseShapeOp::attachInterface<CollapseShapeOpInterface>(*ctx);
- DimOp::attachInterface<DimOpInterface>(*ctx);
- EmptyOp::attachInterface<EmptyOpInterface>(*ctx);
- ExpandShapeOp::attachInterface<ExpandShapeOpInterface>(*ctx);
- ExtractSliceOp::attachInterface<ExtractSliceOpInterface>(*ctx);
- ExtractOp::attachInterface<ExtractOpInterface>(*ctx);
- FromElementsOp::attachInterface<FromElementsOpInterface>(*ctx);
- GenerateOp::attachInterface<GenerateOpInterface>(*ctx);
- InsertOp::attachInterface<InsertOpInterface>(*ctx);
- InsertSliceOp::attachInterface<InsertSliceOpInterface>(*ctx);
- PadOp::attachInterface<PadOpInterface>(*ctx);
- ParallelInsertSliceOp::attachInterface<ParallelInsertSliceOpInterface>(
- *ctx);
- RankOp::attachInterface<RankOpInterface>(*ctx);
- ReshapeOp::attachInterface<ReshapeOpInterface>(*ctx);
- SplatOp::attachInterface<SplatOpInterface>(*ctx);
-
- // Load additional dialects of which ops may get created.
- ctx->loadDialect<arith::ArithDialect, linalg::LinalgDialect>();
- });
+ registry.addExtension(
+ "TENSOR_BUFFERIZATION",
+ +[](MLIRContext *ctx, tensor::TensorDialect *dialect) {
+ CastOp::attachInterface<CastOpInterface>(*ctx);
+ CollapseShapeOp::attachInterface<CollapseShapeOpInterface>(*ctx);
+ DimOp::attachInterface<DimOpInterface>(*ctx);
+ EmptyOp::attachInterface<EmptyOpInterface>(*ctx);
+ ExpandShapeOp::attachInterface<ExpandShapeOpInterface>(*ctx);
+ ExtractSliceOp::attachInterface<ExtractSliceOpInterface>(*ctx);
+ ExtractOp::attachInterface<ExtractOpInterface>(*ctx);
+ FromElementsOp::attachInterface<FromElementsOpInterface>(*ctx);
+ GenerateOp::attachInterface<GenerateOpInterface>(*ctx);
+ InsertOp::attachInterface<InsertOpInterface>(*ctx);
+ InsertSliceOp::attachInterface<InsertSliceOpInterface>(*ctx);
+ PadOp::attachInterface<PadOpInterface>(*ctx);
+ ParallelInsertSliceOp::attachInterface<ParallelInsertSliceOpInterface>(
+ *ctx);
+ RankOp::attachInterface<RankOpInterface>(*ctx);
+ ReshapeOp::attachInterface<ReshapeOpInterface>(*ctx);
+ SplatOp::attachInterface<SplatOpInterface>(*ctx);
+
+ // Load additional dialects of which ops may get created.
+ ctx->loadDialect<arith::ArithDialect, linalg::LinalgDialect>();
+ });
// Bufferization requires SubsetInsertionOpInterface models. Make sure that
// they are registered.
diff --git a/mlir/lib/Dialect/Tensor/Transforms/SubsetInsertionOpInterfaceImpl.cpp b/mlir/lib/Dialect/Tensor/Transforms/SubsetInsertionOpInterfaceImpl.cpp
index d50d7c62b789c..2b052cbb762e1 100644
--- a/mlir/lib/Dialect/Tensor/Transforms/SubsetInsertionOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Tensor/Transforms/SubsetInsertionOpInterfaceImpl.cpp
@@ -86,20 +86,21 @@ struct InsertSliceLikeOpSubsetInsertionOpInterface
void mlir::tensor::registerSubsetOpInterfaceExternalModels(
DialectRegistry ®istry) {
- registry.addExtension(+[](MLIRContext *ctx, tensor::TensorDialect *dialect) {
- // Note: `SubsetExtractionOpInterface` and `SubsetInsertionOpInterface`
- // require `SubsetOpInterface`.
- ExtractSliceOp::attachInterface<ExtractSliceOpSubsetOpInterface>(*ctx);
- ExtractSliceOp::attachInterface<ExtractSliceOpSubsetExtractionOpInterface>(
- *ctx);
- InsertSliceOp::attachInterface<
- InsertSliceLikeOpSubsetOpInterface<InsertSliceOp>>(*ctx);
- InsertSliceOp::attachInterface<
- InsertSliceLikeOpSubsetInsertionOpInterface<InsertSliceOp>>(*ctx);
- ParallelInsertSliceOp::attachInterface<
- InsertSliceLikeOpSubsetOpInterface<ParallelInsertSliceOp>>(*ctx);
- ParallelInsertSliceOp::attachInterface<
- InsertSliceLikeOpSubsetInsertionOpInterface<ParallelInsertSliceOp>>(
- *ctx);
- });
+ registry.addExtension(
+ "TENSOR_SUBSET", +[](MLIRContext *ctx, tensor::TensorDialect *dialect) {
+ // Note: `SubsetExtractionOpInterface` and `SubsetInsertionOpInterface`
+ // require `SubsetOpInterface`.
+ ExtractSliceOp::attachInterface<ExtractSliceOpSubsetOpInterface>(*ctx);
+ ExtractSliceOp::attachInterface<
+ ExtractSliceOpSubsetExtractionOpInterface>(*ctx);
+ InsertSliceOp::attachInterface<
+ InsertSliceLikeOpSubsetOpInterface<InsertSliceOp>>(*ctx);
+ InsertSliceOp::attachInterface<
+ InsertSliceLikeOpSubsetInsertionOpInterface<InsertSliceOp>>(*ctx);
+ ParallelInsertSliceOp::attachInterface<
+ InsertSliceLikeOpSubsetOpInterface<ParallelInsertSliceOp>>(*ctx);
+ ParallelInsertSliceOp::attachInterface<
+ InsertSliceLikeOpSubsetInsertionOpInterface<ParallelInsertSliceOp>>(
+ *ctx);
+ });
}
diff --git a/mlir/lib/Dialect/Tosa/IR/ShardingInterfaceImpl.cpp b/mlir/lib/Dialect/Tosa/IR/ShardingInterfaceImpl.cpp
index ffbb707344b8c..6161d5ae7e3ee 100644
--- a/mlir/lib/Dialect/Tosa/IR/ShardingInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Tosa/IR/ShardingInterfaceImpl.cpp
@@ -76,15 +76,17 @@ static void registerElemwiseAll(MLIRContext *ctx) {
void mlir::tosa::registerShardingInterfaceExternalModels(
DialectRegistry ®istry) {
- registry.addExtension(+[](MLIRContext *ctx, TosaDialect *dialect) {
- registerElemwiseAll<
- ClampOp, SigmoidOp, TanhOp, AddOp, ArithmeticRightShiftOp, BitwiseAndOp,
- BitwiseOrOp, BitwiseXorOp, IntDivOp, LogicalAndOp, LogicalLeftShiftOp,
- LogicalRightShiftOp, LogicalOrOp, LogicalXorOp, MaximumOp, MinimumOp,
- MulOp, PowOp, SubOp, AbsOp, BitwiseNotOp, CeilOp, ClzOp, ExpOp, FloorOp,
- LogOp, LogicalNotOp, NegateOp, ReciprocalOp, RsqrtOp, SelectOp, EqualOp,
- GreaterOp, GreaterEqualOp>(ctx);
+ registry.addExtension(
+ "TOSA_SHARDING", +[](MLIRContext *ctx, TosaDialect *dialect) {
+ registerElemwiseAll<
+ ClampOp, SigmoidOp, TanhOp, AddOp, ArithmeticRightShiftOp,
+ BitwiseAndOp, BitwiseOrOp, BitwiseXorOp, IntDivOp, LogicalAndOp,
+ LogicalLeftShiftOp, LogicalRightShiftOp, LogicalOrOp, LogicalXorOp,
+ MaximumOp, MinimumOp, MulOp, PowOp, SubOp, AbsOp, BitwiseNotOp,
+ CeilOp, ClzOp, ExpOp, FloorOp, LogOp, LogicalNotOp, NegateOp,
+ ReciprocalOp, RsqrtOp, SelectOp, EqualOp, GreaterOp,
+ GreaterEqualOp>(ctx);
- MatMulOp::attachInterface<MatMulOpSharding>(*ctx);
- });
+ MatMulOp::attachInterface<MatMulOpSharding>(*ctx);
+ });
}
diff --git a/mlir/lib/Dialect/Transform/DebugExtension/DebugExtension.cpp b/mlir/lib/Dialect/Transform/DebugExtension/DebugExtension.cpp
index e369daddb00cb..1e3c426162462 100644
--- a/mlir/lib/Dialect/Transform/DebugExtension/DebugExtension.cpp
+++ b/mlir/lib/Dialect/Transform/DebugExtension/DebugExtension.cpp
@@ -26,6 +26,8 @@ class DebugExtension
#include "mlir/Dialect/Transform/DebugExtension/DebugExtensionOps.cpp.inc"
>();
}
+
+ static constexpr llvm::StringRef extensionID = "TRANSFORM_DEBUG";
};
} // namespace
diff --git a/mlir/lib/Dialect/Transform/IRDLExtension/IRDLExtension.cpp b/mlir/lib/Dialect/Transform/IRDLExtension/IRDLExtension.cpp
index 94004365b8a1a..8cae64ff8de0a 100644
--- a/mlir/lib/Dialect/Transform/IRDLExtension/IRDLExtension.cpp
+++ b/mlir/lib/Dialect/Transform/IRDLExtension/IRDLExtension.cpp
@@ -26,6 +26,8 @@ class IRDLExtension
declareDependentDialect<irdl::IRDLDialect>();
}
+
+ static constexpr llvm::StringRef extensionID = "IRDL_TRANSFORM";
};
} // namespace
diff --git a/mlir/lib/Dialect/Transform/LoopExtension/LoopExtension.cpp b/mlir/lib/Dialect/Transform/LoopExtension/LoopExtension.cpp
index b33288fd7b991..032cd65e1065f 100644
--- a/mlir/lib/Dialect/Transform/LoopExtension/LoopExtension.cpp
+++ b/mlir/lib/Dialect/Transform/LoopExtension/LoopExtension.cpp
@@ -26,6 +26,8 @@ class LoopExtension
#include "mlir/Dialect/Transform/LoopExtension/LoopExtensionOps.cpp.inc"
>();
}
+
+ static constexpr llvm::StringRef extensionID = "TRANSFORM_LOOP";
};
} // namespace
diff --git a/mlir/lib/Dialect/Transform/PDLExtension/PDLExtension.cpp b/mlir/lib/Dialect/Transform/PDLExtension/PDLExtension.cpp
index 2c770abd56d52..c7f1e63e3d921 100644
--- a/mlir/lib/Dialect/Transform/PDLExtension/PDLExtension.cpp
+++ b/mlir/lib/Dialect/Transform/PDLExtension/PDLExtension.cpp
@@ -61,6 +61,8 @@ class PDLExtension : public transform::TransformDialectExtension<PDLExtension> {
PDLOperationTypeTransformHandleTypeInterfaceImpl>(*context);
});
}
+
+ static constexpr llvm::StringRef extensionID = "TRANSFORM_PDL";
};
} // namespace
diff --git a/mlir/lib/Dialect/Vector/IR/ValueBoundsOpInterfaceImpl.cpp b/mlir/lib/Dialect/Vector/IR/ValueBoundsOpInterfaceImpl.cpp
index ca95072d9bb0f..91bf4ef633a6d 100644
--- a/mlir/lib/Dialect/Vector/IR/ValueBoundsOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Vector/IR/ValueBoundsOpInterfaceImpl.cpp
@@ -44,8 +44,10 @@ struct VectorScaleOpInterface
void mlir::vector::registerValueBoundsOpInterfaceExternalModels(
DialectRegistry ®istry) {
- registry.addExtension(+[](MLIRContext *ctx, vector::VectorDialect *dialect) {
- vector::VectorScaleOp::attachInterface<vector::VectorScaleOpInterface>(
- *ctx);
- });
+ registry.addExtension(
+ "VECTOR_VALUE_BOUNDS",
+ +[](MLIRContext *ctx, vector::VectorDialect *dialect) {
+ vector::VectorScaleOp::attachInterface<vector::VectorScaleOpInterface>(
+ *ctx);
+ });
}
diff --git a/mlir/lib/Dialect/Vector/TransformOps/VectorTransformOps.cpp b/mlir/lib/Dialect/Vector/TransformOps/VectorTransformOps.cpp
index 2e9aa88011825..0b42040a4d63b 100644
--- a/mlir/lib/Dialect/Vector/TransformOps/VectorTransformOps.cpp
+++ b/mlir/lib/Dialect/Vector/TransformOps/VectorTransformOps.cpp
@@ -220,6 +220,8 @@ class VectorTransformDialectExtension
#include "mlir/Dialect/Vector/TransformOps/VectorTransformOps.cpp.inc"
>();
}
+
+ static constexpr llvm::StringRef extensionID = "VECTOR_TRANSFORM";
};
} // namespace
diff --git a/mlir/lib/Dialect/Vector/Transforms/BufferizableOpInterfaceImpl.cpp b/mlir/lib/Dialect/Vector/Transforms/BufferizableOpInterfaceImpl.cpp
index 1caec5bb8644f..d3e01e33b3f1d 100644
--- a/mlir/lib/Dialect/Vector/Transforms/BufferizableOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/BufferizableOpInterfaceImpl.cpp
@@ -317,11 +317,13 @@ struct YieldOpInterface
void mlir::vector::registerBufferizableOpInterfaceExternalModels(
DialectRegistry ®istry) {
- registry.addExtension(+[](MLIRContext *ctx, vector::VectorDialect *dialect) {
- TransferReadOp::attachInterface<TransferReadOpInterface>(*ctx);
- TransferWriteOp::attachInterface<TransferWriteOpInterface>(*ctx);
- GatherOp::attachInterface<GatherOpInterface>(*ctx);
- MaskOp::attachInterface<MaskOpInterface>(*ctx);
- YieldOp::attachInterface<YieldOpInterface>(*ctx);
- });
+ registry.addExtension(
+ "VECTOR_BUFFERIZATION",
+ +[](MLIRContext *ctx, vector::VectorDialect *dialect) {
+ TransferReadOp::attachInterface<TransferReadOpInterface>(*ctx);
+ TransferWriteOp::attachInterface<TransferWriteOpInterface>(*ctx);
+ GatherOp::attachInterface<GatherOpInterface>(*ctx);
+ MaskOp::attachInterface<MaskOpInterface>(*ctx);
+ YieldOp::attachInterface<YieldOpInterface>(*ctx);
+ });
}
diff --git a/mlir/lib/Dialect/Vector/Transforms/SubsetOpInterfaceImpl.cpp b/mlir/lib/Dialect/Vector/Transforms/SubsetOpInterfaceImpl.cpp
index b450d5b78a466..91a014a7ab24b 100644
--- a/mlir/lib/Dialect/Vector/Transforms/SubsetOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/SubsetOpInterfaceImpl.cpp
@@ -69,14 +69,15 @@ struct TransferWriteOpSubsetInsertionOpInterface
void mlir::vector::registerSubsetOpInterfaceExternalModels(
DialectRegistry ®istry) {
- registry.addExtension(+[](MLIRContext *ctx, vector::VectorDialect *dialect) {
- TransferReadOp::attachInterface<XferOpSubsetOpInterface<TransferReadOp>>(
- *ctx);
- TransferReadOp::attachInterface<TransferReadOpSubsetExtractionOpInterface>(
- *ctx);
- TransferWriteOp::attachInterface<XferOpSubsetOpInterface<TransferWriteOp>>(
- *ctx);
- TransferWriteOp::attachInterface<TransferWriteOpSubsetInsertionOpInterface>(
- *ctx);
- });
+ registry.addExtension(
+ "VECTOR_SUBSET", +[](MLIRContext *ctx, vector::VectorDialect *dialect) {
+ TransferReadOp::attachInterface<
+ XferOpSubsetOpInterface<TransferReadOp>>(*ctx);
+ TransferReadOp::attachInterface<
+ TransferReadOpSubsetExtractionOpInterface>(*ctx);
+ TransferWriteOp::attachInterface<
+ XferOpSubsetOpInterface<TransferWriteOp>>(*ctx);
+ TransferWriteOp::attachInterface<
+ TransferWriteOpSubsetInsertionOpInterface>(*ctx);
+ });
}
diff --git a/mlir/lib/Interfaces/CastInterfaces.cpp b/mlir/lib/Interfaces/CastInterfaces.cpp
index 05c872daf5dab..11e581a5fefa8 100644
--- a/mlir/lib/Interfaces/CastInterfaces.cpp
+++ b/mlir/lib/Interfaces/CastInterfaces.cpp
@@ -79,10 +79,11 @@ struct UnrealizedConversionCastOpInterface
void mlir::builtin::registerCastOpInterfaceExternalModels(
DialectRegistry ®istry) {
- registry.addExtension(+[](MLIRContext *ctx, BuiltinDialect *dialect) {
- UnrealizedConversionCastOp::attachInterface<
- UnrealizedConversionCastOpInterface>(*ctx);
- });
+ registry.addExtension(
+ "BUILTIN_CAST", +[](MLIRContext *ctx, BuiltinDialect *dialect) {
+ UnrealizedConversionCastOp::attachInterface<
+ UnrealizedConversionCastOpInterface>(*ctx);
+ });
}
//===----------------------------------------------------------------------===//
diff --git a/mlir/lib/Target/LLVM/NVVM/Target.cpp b/mlir/lib/Target/LLVM/NVVM/Target.cpp
index e608d26e8d2ec..f5d1babd01180 100644
--- a/mlir/lib/Target/LLVM/NVVM/Target.cpp
+++ b/mlir/lib/Target/LLVM/NVVM/Target.cpp
@@ -58,9 +58,10 @@ class NVVMTargetAttrImpl
// Register the NVVM dialect, the NVVM translation & the target interface.
void mlir::NVVM::registerNVVMTargetInterfaceExternalModels(
DialectRegistry ®istry) {
- registry.addExtension(+[](MLIRContext *ctx, NVVM::NVVMDialect *dialect) {
- NVVMTargetAttr::attachInterface<NVVMTargetAttrImpl>(*ctx);
- });
+ registry.addExtension(
+ "NVVM_TARGET", +[](MLIRContext *ctx, NVVM::NVVMDialect *dialect) {
+ NVVMTargetAttr::attachInterface<NVVMTargetAttrImpl>(*ctx);
+ });
}
void mlir::NVVM::registerNVVMTargetInterfaceExternalModels(
diff --git a/mlir/lib/Target/LLVM/ROCDL/Target.cpp b/mlir/lib/Target/LLVM/ROCDL/Target.cpp
index 4d23f987eb05e..3a56688117fda 100644
--- a/mlir/lib/Target/LLVM/ROCDL/Target.cpp
+++ b/mlir/lib/Target/LLVM/ROCDL/Target.cpp
@@ -68,9 +68,10 @@ class ROCDLTargetAttrImpl
// Register the ROCDL dialect, the ROCDL translation and the target interface.
void mlir::ROCDL::registerROCDLTargetInterfaceExternalModels(
DialectRegistry ®istry) {
- registry.addExtension(+[](MLIRContext *ctx, ROCDL::ROCDLDialect *dialect) {
- ROCDLTargetAttr::attachInterface<ROCDLTargetAttrImpl>(*ctx);
- });
+ registry.addExtension(
+ "ROCDL_TARGET", +[](MLIRContext *ctx, ROCDL::ROCDLDialect *dialect) {
+ ROCDLTargetAttr::attachInterface<ROCDLTargetAttrImpl>(*ctx);
+ });
}
void mlir::ROCDL::registerROCDLTargetInterfaceExternalModels(
diff --git a/mlir/lib/Target/LLVMIR/Dialect/AMX/AMXToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/AMX/AMXToLLVMIRTranslation.cpp
index 044462d33cfd1..4ebcfa5941d0e 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/AMX/AMXToLLVMIRTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/AMX/AMXToLLVMIRTranslation.cpp
@@ -44,9 +44,10 @@ class AMXDialectLLVMIRTranslationInterface
void mlir::registerAMXDialectTranslation(DialectRegistry ®istry) {
registry.insert<amx::AMXDialect>();
- registry.addExtension(+[](MLIRContext *ctx, amx::AMXDialect *dialect) {
- dialect->addInterfaces<AMXDialectLLVMIRTranslationInterface>();
- });
+ registry.addExtension(
+ "AMX_TO_LLVMIR", +[](MLIRContext *ctx, amx::AMXDialect *dialect) {
+ dialect->addInterfaces<AMXDialectLLVMIRTranslationInterface>();
+ });
}
void mlir::registerAMXDialectTranslation(MLIRContext &context) {
diff --git a/mlir/lib/Target/LLVMIR/Dialect/ArmNeon/ArmNeonToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/ArmNeon/ArmNeonToLLVMIRTranslation.cpp
index 7098592d506e0..d317f68731eb8 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/ArmNeon/ArmNeonToLLVMIRTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/ArmNeon/ArmNeonToLLVMIRTranslation.cpp
@@ -46,6 +46,7 @@ class ArmNeonDialectLLVMIRTranslationInterface
void mlir::registerArmNeonDialectTranslation(DialectRegistry ®istry) {
registry.insert<arm_neon::ArmNeonDialect>();
registry.addExtension(
+ "ARM_NEON_TO_LLVMIR",
+[](MLIRContext *ctx, arm_neon::ArmNeonDialect *dialect) {
dialect->addInterfaces<ArmNeonDialectLLVMIRTranslationInterface>();
});
diff --git a/mlir/lib/Target/LLVMIR/Dialect/ArmSME/ArmSMEToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/ArmSME/ArmSMEToLLVMIRTranslation.cpp
index e6ee41188d594..a1f63b9f07bdd 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/ArmSME/ArmSMEToLLVMIRTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/ArmSME/ArmSMEToLLVMIRTranslation.cpp
@@ -45,9 +45,11 @@ class ArmSMEDialectLLVMIRTranslationInterface
void mlir::registerArmSMEDialectTranslation(DialectRegistry ®istry) {
registry.insert<arm_sme::ArmSMEDialect>();
- registry.addExtension(+[](MLIRContext *ctx, arm_sme::ArmSMEDialect *dialect) {
- dialect->addInterfaces<ArmSMEDialectLLVMIRTranslationInterface>();
- });
+ registry.addExtension(
+ "ARM_SME_TO_LLVMIR",
+ +[](MLIRContext *ctx, arm_sme::ArmSMEDialect *dialect) {
+ dialect->addInterfaces<ArmSMEDialectLLVMIRTranslationInterface>();
+ });
}
void mlir::registerArmSMEDialectTranslation(MLIRContext &context) {
diff --git a/mlir/lib/Target/LLVMIR/Dialect/ArmSVE/ArmSVEToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/ArmSVE/ArmSVEToLLVMIRTranslation.cpp
index cd10811b68f02..3411f5d147a34 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/ArmSVE/ArmSVEToLLVMIRTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/ArmSVE/ArmSVEToLLVMIRTranslation.cpp
@@ -44,9 +44,11 @@ class ArmSVEDialectLLVMIRTranslationInterface
void mlir::registerArmSVEDialectTranslation(DialectRegistry ®istry) {
registry.insert<arm_sve::ArmSVEDialect>();
- registry.addExtension(+[](MLIRContext *ctx, arm_sve::ArmSVEDialect *dialect) {
- dialect->addInterfaces<ArmSVEDialectLLVMIRTranslationInterface>();
- });
+ registry.addExtension(
+ "ARM_SVE_TO_LLVMIR",
+ +[](MLIRContext *ctx, arm_sve::ArmSVEDialect *dialect) {
+ dialect->addInterfaces<ArmSVEDialectLLVMIRTranslationInterface>();
+ });
}
void mlir::registerArmSVEDialectTranslation(MLIRContext &context) {
diff --git a/mlir/lib/Target/LLVMIR/Dialect/Builtin/BuiltinToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/Builtin/BuiltinToLLVMIRTranslation.cpp
index 51c304cfbb8e5..c120eb243b783 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/Builtin/BuiltinToLLVMIRTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/Builtin/BuiltinToLLVMIRTranslation.cpp
@@ -34,9 +34,10 @@ class BuiltinDialectLLVMIRTranslationInterface
} // namespace
void mlir::registerBuiltinDialectTranslation(DialectRegistry ®istry) {
- registry.addExtension(+[](MLIRContext *ctx, BuiltinDialect *dialect) {
- dialect->addInterfaces<BuiltinDialectLLVMIRTranslationInterface>();
- });
+ registry.addExtension(
+ "BUILTIN_TO_LLVMIR", +[](MLIRContext *ctx, BuiltinDialect *dialect) {
+ dialect->addInterfaces<BuiltinDialectLLVMIRTranslationInterface>();
+ });
}
void mlir::registerBuiltinDialectTranslation(MLIRContext &context) {
diff --git a/mlir/lib/Target/LLVMIR/Dialect/GPU/GPUToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/GPU/GPUToLLVMIRTranslation.cpp
index eecc8f1001ca4..c4fff8cfee7d1 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/GPU/GPUToLLVMIRTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/GPU/GPUToLLVMIRTranslation.cpp
@@ -65,9 +65,10 @@ class GPUDialectLLVMIRTranslationInterface
void mlir::registerGPUDialectTranslation(DialectRegistry ®istry) {
registry.insert<gpu::GPUDialect>();
- registry.addExtension(+[](MLIRContext *ctx, gpu::GPUDialect *dialect) {
- dialect->addInterfaces<GPUDialectLLVMIRTranslationInterface>();
- });
+ registry.addExtension(
+ "GPU_TO_LLVMIR", +[](MLIRContext *ctx, gpu::GPUDialect *dialect) {
+ dialect->addInterfaces<GPUDialectLLVMIRTranslationInterface>();
+ });
}
void mlir::registerGPUDialectTranslation(MLIRContext &context) {
diff --git a/mlir/lib/Target/LLVMIR/Dialect/GPU/SelectObjectAttr.cpp b/mlir/lib/Target/LLVMIR/Dialect/GPU/SelectObjectAttr.cpp
index b023c4c126da3..c982e5f67900f 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/GPU/SelectObjectAttr.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/GPU/SelectObjectAttr.cpp
@@ -56,9 +56,10 @@ std::string getBinaryIdentifier(StringRef binaryName) {
void mlir::gpu::registerOffloadingLLVMTranslationInterfaceExternalModels(
DialectRegistry ®istry) {
- registry.addExtension(+[](MLIRContext *ctx, gpu::GPUDialect *dialect) {
- SelectObjectAttr::attachInterface<SelectObjectAttrImpl>(*ctx);
- });
+ registry.addExtension(
+ "GPU_OFFLOADING_LLVM", +[](MLIRContext *ctx, gpu::GPUDialect *dialect) {
+ SelectObjectAttr::attachInterface<SelectObjectAttrImpl>(*ctx);
+ });
}
gpu::ObjectAttr
diff --git a/mlir/lib/Target/LLVMIR/Dialect/LLVMIR/LLVMIRToLLVMTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/LLVMIR/LLVMIRToLLVMTranslation.cpp
index 06673965245c0..14485c54ffdde 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/LLVMIR/LLVMIRToLLVMTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/LLVMIR/LLVMIRToLLVMTranslation.cpp
@@ -281,9 +281,10 @@ class LLVMDialectLLVMIRImportInterface : public LLVMImportDialectInterface {
void mlir::registerLLVMDialectImport(DialectRegistry ®istry) {
registry.insert<LLVM::LLVMDialect>();
- registry.addExtension(+[](MLIRContext *ctx, LLVM::LLVMDialect *dialect) {
- dialect->addInterfaces<LLVMDialectLLVMIRImportInterface>();
- });
+ registry.addExtension(
+ "LLVMIR_TO_LLVM", +[](MLIRContext *ctx, LLVM::LLVMDialect *dialect) {
+ dialect->addInterfaces<LLVMDialectLLVMIRImportInterface>();
+ });
}
void mlir::registerLLVMDialectImport(MLIRContext &context) {
diff --git a/mlir/lib/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.cpp
index bdb15a290209b..d5bc600fd7e13 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.cpp
@@ -439,9 +439,10 @@ class LLVMDialectLLVMIRTranslationInterface
void mlir::registerLLVMDialectTranslation(DialectRegistry ®istry) {
registry.insert<LLVM::LLVMDialect>();
- registry.addExtension(+[](MLIRContext *ctx, LLVM::LLVMDialect *dialect) {
- dialect->addInterfaces<LLVMDialectLLVMIRTranslationInterface>();
- });
+ registry.addExtension(
+ "LLVM_TO_LLVMIR", +[](MLIRContext *ctx, LLVM::LLVMDialect *dialect) {
+ dialect->addInterfaces<LLVMDialectLLVMIRTranslationInterface>();
+ });
}
void mlir::registerLLVMDialectTranslation(MLIRContext &context) {
diff --git a/mlir/lib/Target/LLVMIR/Dialect/NVVM/LLVMIRToNVVMTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/NVVM/LLVMIRToNVVMTranslation.cpp
index 855abc12a909e..9978513e59ec0 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/NVVM/LLVMIRToNVVMTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/NVVM/LLVMIRToNVVMTranslation.cpp
@@ -81,9 +81,10 @@ class NVVMDialectLLVMIRImportInterface : public LLVMImportDialectInterface {
void mlir::registerNVVMDialectImport(DialectRegistry ®istry) {
registry.insert<NVVM::NVVMDialect>();
- registry.addExtension(+[](MLIRContext *ctx, NVVM::NVVMDialect *dialect) {
- dialect->addInterfaces<NVVMDialectLLVMIRImportInterface>();
- });
+ registry.addExtension(
+ "LLVMIR_TO_NVVM", +[](MLIRContext *ctx, NVVM::NVVMDialect *dialect) {
+ dialect->addInterfaces<NVVMDialectLLVMIRImportInterface>();
+ });
}
void mlir::registerNVVMDialectImport(MLIRContext &context) {
diff --git a/mlir/lib/Target/LLVMIR/Dialect/NVVM/NVVMToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/NVVM/NVVMToLLVMIRTranslation.cpp
index a09c24dda82af..e12364b1122f0 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/NVVM/NVVMToLLVMIRTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/NVVM/NVVMToLLVMIRTranslation.cpp
@@ -261,9 +261,10 @@ class NVVMDialectLLVMIRTranslationInterface
void mlir::registerNVVMDialectTranslation(DialectRegistry ®istry) {
registry.insert<NVVM::NVVMDialect>();
- registry.addExtension(+[](MLIRContext *ctx, NVVM::NVVMDialect *dialect) {
- dialect->addInterfaces<NVVMDialectLLVMIRTranslationInterface>();
- });
+ registry.addExtension(
+ "NVVM_TO_LLVMIR", +[](MLIRContext *ctx, NVVM::NVVMDialect *dialect) {
+ dialect->addInterfaces<NVVMDialectLLVMIRTranslationInterface>();
+ });
}
void mlir::registerNVVMDialectTranslation(MLIRContext &context) {
diff --git a/mlir/lib/Target/LLVMIR/Dialect/OpenACC/OpenACCToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/OpenACC/OpenACCToLLVMIRTranslation.cpp
index d9cf85e4aecab..67f0f8fc4f374 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/OpenACC/OpenACCToLLVMIRTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/OpenACC/OpenACCToLLVMIRTranslation.cpp
@@ -532,9 +532,10 @@ LogicalResult OpenACCDialectLLVMIRTranslationInterface::convertOperation(
void mlir::registerOpenACCDialectTranslation(DialectRegistry ®istry) {
registry.insert<acc::OpenACCDialect>();
- registry.addExtension(+[](MLIRContext *ctx, acc::OpenACCDialect *dialect) {
- dialect->addInterfaces<OpenACCDialectLLVMIRTranslationInterface>();
- });
+ registry.addExtension(
+ "OPENACC_TO_LLVMIR", +[](MLIRContext *ctx, acc::OpenACCDialect *dialect) {
+ dialect->addInterfaces<OpenACCDialectLLVMIRTranslationInterface>();
+ });
}
void mlir::registerOpenACCDialectTranslation(MLIRContext &context) {
diff --git a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
index ddee117838697..a43761460fa31 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
@@ -3741,9 +3741,10 @@ LogicalResult OpenMPDialectLLVMIRTranslationInterface::convertOperation(
void mlir::registerOpenMPDialectTranslation(DialectRegistry ®istry) {
registry.insert<omp::OpenMPDialect>();
- registry.addExtension(+[](MLIRContext *ctx, omp::OpenMPDialect *dialect) {
- dialect->addInterfaces<OpenMPDialectLLVMIRTranslationInterface>();
- });
+ registry.addExtension(
+ "OPENMV_TO_LLVMIR", +[](MLIRContext *ctx, omp::OpenMPDialect *dialect) {
+ dialect->addInterfaces<OpenMPDialectLLVMIRTranslationInterface>();
+ });
}
void mlir::registerOpenMPDialectTranslation(MLIRContext &context) {
diff --git a/mlir/lib/Target/LLVMIR/Dialect/ROCDL/ROCDLToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/ROCDL/ROCDLToLLVMIRTranslation.cpp
index 2a146f5efed30..7fe58e5d161a8 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/ROCDL/ROCDLToLLVMIRTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/ROCDL/ROCDLToLLVMIRTranslation.cpp
@@ -190,9 +190,10 @@ class ROCDLDialectLLVMIRTranslationInterface
void mlir::registerROCDLDialectTranslation(DialectRegistry ®istry) {
registry.insert<ROCDL::ROCDLDialect>();
- registry.addExtension(+[](MLIRContext *ctx, ROCDL::ROCDLDialect *dialect) {
- dialect->addInterfaces<ROCDLDialectLLVMIRTranslationInterface>();
- });
+ registry.addExtension(
+ "ROCDL_TO_LLVMIR", +[](MLIRContext *ctx, ROCDL::ROCDLDialect *dialect) {
+ dialect->addInterfaces<ROCDLDialectLLVMIRTranslationInterface>();
+ });
}
void mlir::registerROCDLDialectTranslation(MLIRContext &context) {
diff --git a/mlir/lib/Target/LLVMIR/Dialect/VCIX/VCIXToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/VCIX/VCIXToLLVMIRTranslation.cpp
index b78b002d32292..17d63863d9ddc 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/VCIX/VCIXToLLVMIRTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/VCIX/VCIXToLLVMIRTranslation.cpp
@@ -76,9 +76,10 @@ class VCIXDialectLLVMIRTranslationInterface
void mlir::registerVCIXDialectTranslation(DialectRegistry ®istry) {
registry.insert<vcix::VCIXDialect>();
- registry.addExtension(+[](MLIRContext *ctx, vcix::VCIXDialect *dialect) {
- dialect->addInterfaces<VCIXDialectLLVMIRTranslationInterface>();
- });
+ registry.addExtension(
+ "VCIX_TO_LLVMIR", +[](MLIRContext *ctx, vcix::VCIXDialect *dialect) {
+ dialect->addInterfaces<VCIXDialectLLVMIRTranslationInterface>();
+ });
}
void mlir::registerVCIXDialectTranslation(MLIRContext &context) {
diff --git a/mlir/lib/Target/LLVMIR/Dialect/X86Vector/X86VectorToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/X86Vector/X86VectorToLLVMIRTranslation.cpp
index fa5f61420ee8a..5a6104162123b 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/X86Vector/X86VectorToLLVMIRTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/X86Vector/X86VectorToLLVMIRTranslation.cpp
@@ -46,6 +46,7 @@ class X86VectorDialectLLVMIRTranslationInterface
void mlir::registerX86VectorDialectTranslation(DialectRegistry ®istry) {
registry.insert<x86vector::X86VectorDialect>();
registry.addExtension(
+ "X86_TO_LLVMIR",
+[](MLIRContext *ctx, x86vector::X86VectorDialect *dialect) {
dialect->addInterfaces<X86VectorDialectLLVMIRTranslationInterface>();
});
diff --git a/mlir/lib/Target/SPIRV/Target.cpp b/mlir/lib/Target/SPIRV/Target.cpp
index 4c416abe71cac..15832a43b6620 100644
--- a/mlir/lib/Target/SPIRV/Target.cpp
+++ b/mlir/lib/Target/SPIRV/Target.cpp
@@ -43,9 +43,10 @@ class SPIRVTargetAttrImpl
// Register the SPIR-V dialect, the SPIR-V translation & the target interface.
void mlir::spirv::registerSPIRVTargetInterfaceExternalModels(
DialectRegistry ®istry) {
- registry.addExtension(+[](MLIRContext *ctx, spirv::SPIRVDialect *dialect) {
- spirv::TargetEnvAttr::attachInterface<SPIRVTargetAttrImpl>(*ctx);
- });
+ registry.addExtension(
+ "SPIRV_TARGET", +[](MLIRContext *ctx, spirv::SPIRVDialect *dialect) {
+ spirv::TargetEnvAttr::attachInterface<SPIRVTargetAttrImpl>(*ctx);
+ });
}
void mlir::spirv::registerSPIRVTargetInterfaceExternalModels(
diff --git a/mlir/test/lib/Dialect/Test/TestFromLLVMIRTranslation.cpp b/mlir/test/lib/Dialect/Test/TestFromLLVMIRTranslation.cpp
index dc6413b25707e..0422705663b81 100644
--- a/mlir/test/lib/Dialect/Test/TestFromLLVMIRTranslation.cpp
+++ b/mlir/test/lib/Dialect/Test/TestFromLLVMIRTranslation.cpp
@@ -104,6 +104,7 @@ void registerTestFromLLVMIR() {
registry.insert<test::TestDialect>();
registerLLVMDialectImport(registry);
registry.addExtension(
+ "TEST_FROM_LLVMIR",
+[](MLIRContext *ctx, test::TestDialect *dialect) {
dialect->addInterfaces<TestDialectLLVMImportDialectInterface>();
});
diff --git a/mlir/test/lib/Dialect/Test/TestToLLVMIRTranslation.cpp b/mlir/test/lib/Dialect/Test/TestToLLVMIRTranslation.cpp
index 157c6265be834..3d4b8dc6b1495 100644
--- a/mlir/test/lib/Dialect/Test/TestToLLVMIRTranslation.cpp
+++ b/mlir/test/lib/Dialect/Test/TestToLLVMIRTranslation.cpp
@@ -140,6 +140,7 @@ void registerTestToLLVMIR() {
registerBuiltinDialectTranslation(registry);
registerLLVMDialectTranslation(registry);
registry.addExtension(
+ "TEST_TO_LLVMIR",
+[](MLIRContext *ctx, test::TestDialect *dialect) {
dialect->addInterfaces<TestDialectLLVMIRTranslationInterface>();
});
diff --git a/mlir/test/lib/Dialect/Transform/TestTransformDialectExtension.cpp b/mlir/test/lib/Dialect/Transform/TestTransformDialectExtension.cpp
index b8a4b9470d736..cc8f0e8988ea6 100644
--- a/mlir/test/lib/Dialect/Transform/TestTransformDialectExtension.cpp
+++ b/mlir/test/lib/Dialect/Transform/TestTransformDialectExtension.cpp
@@ -905,6 +905,8 @@ class TestTransformDialectExtension
hooks.mergeInPDLMatchHooks(std::move(constraints));
});
}
+
+ static constexpr llvm::StringRef extensionID = "TEST_TRANSFORM";
};
} // namespace
diff --git a/mlir/test/lib/Interfaces/TilingInterface/TestTilingInterfaceTransformOps.cpp b/mlir/test/lib/Interfaces/TilingInterface/TestTilingInterfaceTransformOps.cpp
index 8f206d9077272..172001123d4da 100644
--- a/mlir/test/lib/Interfaces/TilingInterface/TestTilingInterfaceTransformOps.cpp
+++ b/mlir/test/lib/Interfaces/TilingInterface/TestTilingInterfaceTransformOps.cpp
@@ -399,6 +399,8 @@ class TestTilingInterfaceDialectExtension
#include "TestTilingInterfaceTransformOps.cpp.inc"
>();
}
+
+ static constexpr llvm::StringRef extensionID = "TEST_TILING_TRANSFORM";
};
} // namespace
diff --git a/mlir/unittests/Dialect/Transform/BuildOnlyExtensionTest.cpp b/mlir/unittests/Dialect/Transform/BuildOnlyExtensionTest.cpp
index 40fb752ffd6eb..2d8fbf078618f 100644
--- a/mlir/unittests/Dialect/Transform/BuildOnlyExtensionTest.cpp
+++ b/mlir/unittests/Dialect/Transform/BuildOnlyExtensionTest.cpp
@@ -20,6 +20,8 @@ class Extension : public TransformDialectExtension<Extension> {
public:
using Base::Base;
void init() { declareGeneratedDialect<func::FuncDialect>(); }
+
+ static constexpr llvm::StringRef extensionID = "TRANSFORM_BUILD_ONLY";
};
} // end namespace
diff --git a/mlir/unittests/IR/DialectTest.cpp b/mlir/unittests/IR/DialectTest.cpp
index e99d46e6d2643..b79d7941a12a0 100644
--- a/mlir/unittests/IR/DialectTest.cpp
+++ b/mlir/unittests/IR/DialectTest.cpp
@@ -75,7 +75,8 @@ TEST(Dialect, DelayedInterfaceRegistration) {
registry.insert<TestDialect, SecondTestDialect>();
// Delayed registration of an interface for TestDialect.
- registry.addExtension(+[](MLIRContext *ctx, TestDialect *dialect) {
+ registry.addExtension(
+ "TEST_DIALECT_DELAYED", +[](MLIRContext *ctx, TestDialect *dialect) {
dialect->addInterfaces<TestDialectInterface>();
});
@@ -100,7 +101,7 @@ TEST(Dialect, DelayedInterfaceRegistration) {
DialectRegistry secondRegistry;
secondRegistry.insert<SecondTestDialect>();
secondRegistry.addExtension(
- +[](MLIRContext *ctx, SecondTestDialect *dialect) {
+ "SECOND_TEST", +[](MLIRContext *ctx, SecondTestDialect *dialect) {
dialect->addInterfaces<SecondTestDialectInterface>();
});
context.appendDialectRegistry(secondRegistry);
@@ -113,7 +114,8 @@ TEST(Dialect, RepeatedDelayedRegistration) {
// Set up the delayed registration.
DialectRegistry registry;
registry.insert<TestDialect>();
- registry.addExtension(+[](MLIRContext *ctx, TestDialect *dialect) {
+ registry.addExtension(
+ "TEST_DIALECT", +[](MLIRContext *ctx, TestDialect *dialect) {
dialect->addInterfaces<TestDialectInterface>();
});
MLIRContext context(registry);
@@ -128,7 +130,8 @@ TEST(Dialect, RepeatedDelayedRegistration) {
// on repeated interface registration.
DialectRegistry secondRegistry;
secondRegistry.insert<TestDialect>();
- secondRegistry.addExtension(+[](MLIRContext *ctx, TestDialect *dialect) {
+ secondRegistry.addExtension(
+ "TEST_DIALECT", +[](MLIRContext *ctx, TestDialect *dialect) {
dialect->addInterfaces<TestDialectInterface>();
});
context.appendDialectRegistry(secondRegistry);
@@ -143,12 +146,18 @@ struct DummyExtension : DialectExtension<DummyExtension, TestDialect> {
DummyExtension(int *counter, int numRecursive)
: DialectExtension(), counter(counter), numRecursive(numRecursive) {}
+ inline static std::vector<std::string> extensionIDs;
+
void apply(MLIRContext *ctx, TestDialect *dialect) const final {
++(*counter);
DialectRegistry nestedRegistry;
- for (int i = 0; i < numRecursive; ++i)
+ extensionIDs.reserve(extensionIDs.size() + numRecursive);
+ for (int i = 0; i < numRecursive; ++i) {
+ extensionIDs.push_back("DUMMY_" + std::to_string(i));
nestedRegistry.addExtension(
+ extensionIDs.back(),
std::make_unique<DummyExtension>(counter, /*numRecursive=*/0));
+ }
// Adding additional extensions may trigger a reallocation of the
// `extensions` vector in the dialect registry.
ctx->appendDialectRegistry(nestedRegistry);
@@ -166,10 +175,12 @@ TEST(Dialect, NestedDialectExtension) {
// Add an extension that adds 100 more extensions.
int counter1 = 0;
- registry.addExtension(std::make_unique<DummyExtension>(&counter1, 100));
+ registry.addExtension("DUMMY",
+ std::make_unique<DummyExtension>(&counter1, 100));
// Add one more extension. This should not crash.
int counter2 = 0;
- registry.addExtension(std::make_unique<DummyExtension>(&counter2, 0));
+ registry.addExtension("DUMMY2",
+ std::make_unique<DummyExtension>(&counter2, 0));
// Load dialect and apply extensions.
MLIRContext context(registry);
diff --git a/mlir/unittests/IR/InterfaceAttachmentTest.cpp b/mlir/unittests/IR/InterfaceAttachmentTest.cpp
index b6066dd5685dc..a5b835153b636 100644
--- a/mlir/unittests/IR/InterfaceAttachmentTest.cpp
+++ b/mlir/unittests/IR/InterfaceAttachmentTest.cpp
@@ -103,9 +103,11 @@ TEST(InterfaceAttachment, TypeDelayedContextConstruct) {
// Put the interface in the registry.
DialectRegistry registry;
registry.insert<test::TestDialect>();
- registry.addExtension(+[](MLIRContext *ctx, test::TestDialect *dialect) {
- test::TestType::attachInterface<TestTypeModel>(*ctx);
- });
+ registry.addExtension(
+ "TYPE_DELAYED_CONSTRUCT",
+ +[](MLIRContext *ctx, test::TestDialect *dialect) {
+ test::TestType::attachInterface<TestTypeModel>(*ctx);
+ });
// Check that when a context is constructed with the given registry, the type
// interface gets registered.
@@ -122,9 +124,10 @@ TEST(InterfaceAttachment, TypeDelayedContextAppend) {
// Put the interface in the registry.
DialectRegistry registry;
registry.insert<test::TestDialect>();
- registry.addExtension(+[](MLIRContext *ctx, test::TestDialect *dialect) {
- test::TestType::attachInterface<TestTypeModel>(*ctx);
- });
+ registry.addExtension(
+ "TYPE_DELAYED_APPEND", +[](MLIRContext *ctx, test::TestDialect *dialect) {
+ test::TestType::attachInterface<TestTypeModel>(*ctx);
+ });
// Check that when the registry gets appended to the context, the interface
// becomes available for objects in loaded dialects.
@@ -138,9 +141,10 @@ TEST(InterfaceAttachment, TypeDelayedContextAppend) {
TEST(InterfaceAttachment, RepeatedRegistration) {
DialectRegistry registry;
- registry.addExtension(+[](MLIRContext *ctx, BuiltinDialect *dialect) {
- IntegerType::attachInterface<Model>(*ctx);
- });
+ registry.addExtension(
+ "REPEATED", +[](MLIRContext *ctx, BuiltinDialect *dialect) {
+ IntegerType::attachInterface<Model>(*ctx);
+ });
MLIRContext context(registry);
// Should't fail on repeated registration through the dialect registry.
@@ -151,9 +155,10 @@ TEST(InterfaceAttachment, TypeBuiltinDelayed) {
// Builtin dialect needs to registration or loading, but delayed interface
// registration must still work.
DialectRegistry registry;
- registry.addExtension(+[](MLIRContext *ctx, BuiltinDialect *dialect) {
- IntegerType::attachInterface<Model>(*ctx);
- });
+ registry.addExtension(
+ "BUILTIN_DELAYED", +[](MLIRContext *ctx, BuiltinDialect *dialect) {
+ IntegerType::attachInterface<Model>(*ctx);
+ });
MLIRContext context(registry);
IntegerType i16 = IntegerType::get(&context, 16);
@@ -246,9 +251,10 @@ TEST(InterfaceAttachmentTest, AttributeDelayed) {
// that the delayed registration work for attributes.
DialectRegistry registry;
registry.insert<test::TestDialect>();
- registry.addExtension(+[](MLIRContext *ctx, test::TestDialect *dialect) {
- test::SimpleAAttr::attachInterface<TestExternalSimpleAAttrModel>(*ctx);
- });
+ registry.addExtension(
+ "ATTRIBUTE_DELAYED", +[](MLIRContext *ctx, test::TestDialect *dialect) {
+ test::SimpleAAttr::attachInterface<TestExternalSimpleAAttrModel>(*ctx);
+ });
MLIRContext context(registry);
context.loadDialect<test::TestDialect>();
@@ -352,13 +358,17 @@ struct TestExternalTestOpModel
TEST(InterfaceAttachment, OperationDelayedContextConstruct) {
DialectRegistry registry;
registry.insert<test::TestDialect>();
- registry.addExtension(+[](MLIRContext *ctx, BuiltinDialect *dialect) {
- ModuleOp::attachInterface<TestExternalOpModel>(*ctx);
- });
- registry.addExtension(+[](MLIRContext *ctx, test::TestDialect *dialect) {
- test::OpJ::attachInterface<TestExternalTestOpModel<test::OpJ>>(*ctx);
- test::OpH::attachInterface<TestExternalTestOpModel<test::OpH>>(*ctx);
- });
+ registry.addExtension(
+ "OPERATION_DELAYED_BUILTIN",
+ +[](MLIRContext *ctx, BuiltinDialect *dialect) {
+ ModuleOp::attachInterface<TestExternalOpModel>(*ctx);
+ });
+ registry.addExtension(
+ "OPERATION_DELAYED_TEST",
+ +[](MLIRContext *ctx, test::TestDialect *dialect) {
+ test::OpJ::attachInterface<TestExternalTestOpModel<test::OpJ>>(*ctx);
+ test::OpH::attachInterface<TestExternalTestOpModel<test::OpH>>(*ctx);
+ });
// Construct the context directly from a registry. The interfaces are
// expected to be readily available on operations.
@@ -383,13 +393,17 @@ TEST(InterfaceAttachment, OperationDelayedContextConstruct) {
TEST(InterfaceAttachment, OperationDelayedContextAppend) {
DialectRegistry registry;
registry.insert<test::TestDialect>();
- registry.addExtension(+[](MLIRContext *ctx, BuiltinDialect *dialect) {
- ModuleOp::attachInterface<TestExternalOpModel>(*ctx);
- });
- registry.addExtension(+[](MLIRContext *ctx, test::TestDialect *dialect) {
- test::OpJ::attachInterface<TestExternalTestOpModel<test::OpJ>>(*ctx);
- test::OpH::attachInterface<TestExternalTestOpModel<test::OpH>>(*ctx);
- });
+ registry.addExtension(
+ "OPERATION_DELAYED_BUILTIN",
+ +[](MLIRContext *ctx, BuiltinDialect *dialect) {
+ ModuleOp::attachInterface<TestExternalOpModel>(*ctx);
+ });
+ registry.addExtension(
+ "OPERATION_DELAYED_TEST",
+ +[](MLIRContext *ctx, test::TestDialect *dialect) {
+ test::OpJ::attachInterface<TestExternalTestOpModel<test::OpJ>>(*ctx);
+ test::OpH::attachInterface<TestExternalTestOpModel<test::OpH>>(*ctx);
+ });
// Construct the context, create ops, and only then append the registry. The
// interfaces are expected to be available after appending the registry.
>From 0cf68e846b20027c7903551c6a0d2321f4c06f88 Mon Sep 17 00:00:00 2001
From: Nikhil Kalra <nkalra at apple.com>
Date: Wed, 31 Jul 2024 19:36:45 -0700
Subject: [PATCH 3/3] add dialect registry extension comparison by key
---
mlir/lib/IR/Dialect.cpp | 14 ++++++++++--
mlir/unittests/IR/DialectTest.cpp | 37 ++++++++++++++++++++++++++-----
2 files changed, 43 insertions(+), 8 deletions(-)
diff --git a/mlir/lib/IR/Dialect.cpp b/mlir/lib/IR/Dialect.cpp
index 0bdc34bbc5ae8..c8b55bf751091 100644
--- a/mlir/lib/IR/Dialect.cpp
+++ b/mlir/lib/IR/Dialect.cpp
@@ -291,9 +291,19 @@ void DialectRegistry::applyExtensions(MLIRContext *ctx) const {
}
bool DialectRegistry::isSubsetOf(const DialectRegistry &rhs) const {
- // Treat any extensions conservatively.
- if (!extensions.empty())
+ // Check that all extension keys are present in 'rhs'.
+ llvm::DenseSet<llvm::StringRef> rhsExtensionKeys;
+ {
+ auto rhsKeys = llvm::map_range(rhs.extensions,
+ [](const auto &item) { return item.first; });
+ rhsExtensionKeys.insert(rhsKeys.begin(), rhsKeys.end());
+ }
+
+ if (!llvm::all_of(extensions, [&rhsExtensionKeys](const auto &extension) {
+ return rhsExtensionKeys.contains(extension.first);
+ }))
return false;
+
// Check that the current dialects fully overlap with the dialects in 'rhs'.
return llvm::all_of(
registry, [&](const auto &it) { return rhs.registry.count(it.first); });
diff --git a/mlir/unittests/IR/DialectTest.cpp b/mlir/unittests/IR/DialectTest.cpp
index b79d7941a12a0..2623978639674 100644
--- a/mlir/unittests/IR/DialectTest.cpp
+++ b/mlir/unittests/IR/DialectTest.cpp
@@ -77,8 +77,8 @@ TEST(Dialect, DelayedInterfaceRegistration) {
// Delayed registration of an interface for TestDialect.
registry.addExtension(
"TEST_DIALECT_DELAYED", +[](MLIRContext *ctx, TestDialect *dialect) {
- dialect->addInterfaces<TestDialectInterface>();
- });
+ dialect->addInterfaces<TestDialectInterface>();
+ });
MLIRContext context(registry);
@@ -116,8 +116,8 @@ TEST(Dialect, RepeatedDelayedRegistration) {
registry.insert<TestDialect>();
registry.addExtension(
"TEST_DIALECT", +[](MLIRContext *ctx, TestDialect *dialect) {
- dialect->addInterfaces<TestDialectInterface>();
- });
+ dialect->addInterfaces<TestDialectInterface>();
+ });
MLIRContext context(registry);
// Load the TestDialect and check that the interface got registered for it.
@@ -132,8 +132,8 @@ TEST(Dialect, RepeatedDelayedRegistration) {
secondRegistry.insert<TestDialect>();
secondRegistry.addExtension(
"TEST_DIALECT", +[](MLIRContext *ctx, TestDialect *dialect) {
- dialect->addInterfaces<TestDialectInterface>();
- });
+ dialect->addInterfaces<TestDialectInterface>();
+ });
context.appendDialectRegistry(secondRegistry);
testDialectInterface = dyn_cast<TestDialectInterfaceBase>(testDialect);
EXPECT_TRUE(testDialectInterface != nullptr);
@@ -193,4 +193,29 @@ TEST(Dialect, NestedDialectExtension) {
EXPECT_GE(counter2, 1);
}
+TEST(Dialect, SubsetWithExtensions) {
+ DialectRegistry registry1, registry2;
+ registry1.insert<TestDialect>();
+ registry2.insert<TestDialect>();
+
+ // Validate that the registries are equivalent.
+ ASSERT_TRUE(registry1.isSubsetOf(registry2));
+ ASSERT_TRUE(registry2.isSubsetOf(registry1));
+
+ // Add extensions to registry2.
+ int counter;
+ registry2.addExtension("EXT", std::make_unique<DummyExtension>(&counter, 0));
+
+ // Expect that (1) is a subset of (2) but not the other way around.
+ ASSERT_TRUE(registry1.isSubsetOf(registry2));
+ ASSERT_FALSE(registry2.isSubsetOf(registry1));
+
+ // Add extensions to registry1.
+ registry1.addExtension("EXT", std::make_unique<DummyExtension>(&counter, 0));
+
+ // Expect that (1) and (2) are equivalent.
+ ASSERT_TRUE(registry1.isSubsetOf(registry2));
+ ASSERT_TRUE(registry2.isSubsetOf(registry1));
+}
+
} // namespace
More information about the Mlir-commits
mailing list