[Mlir-commits] [mlir] 7c4e8c6 - [mlir] Disentangle dialect and extension registrations.

Fabian Mora llvmlistbot at llvm.org
Mon Aug 21 17:40:17 PDT 2023 

Author: Nicolas Vasilache
Date: 2023-08-22T00:40:09Z
New Revision: 7c4e8c6a273f25b3ef33e9c123b3969632ab59bb

URL: https://github.com/llvm/llvm-project/commit/7c4e8c6a273f25b3ef33e9c123b3969632ab59bb
DIFF: https://github.com/llvm/llvm-project/commit/7c4e8c6a273f25b3ef33e9c123b3969632ab59bb.diff

LOG: [mlir] Disentangle dialect and extension registrations.

This revision avoids the registration of dialect extensions in Pass::getDependentDialects.

Such registration of extensions can be dangerous because `DialectRegistry::isSubsetOf` is
always guaranteed to return false for extensions (i.e. there is no mechanism to track
whether a lambda is already in the list of already registered extensions).
When the context is already in a multi-threaded mode, this is guaranteed to assert.

Arguably a more structured registration mechanism for extensions with a unique ExtensionID
could be envisioned in the future.

In the process of cleaning this up, multiple usage inconsistencies surfaced around the
registration of translation extensions that this revision also cleans up.

Reviewed By: springerm

Differential Revision: https://reviews.llvm.org/D157703

Added: 
    

Modified: 
    mlir/examples/transform/Ch2/transform-opt/transform-opt.cpp
    mlir/examples/transform/Ch3/transform-opt/transform-opt.cpp
    mlir/include/mlir/Dialect/GPU/Transforms/Passes.h
    mlir/include/mlir/InitAllDialects.h
    mlir/include/mlir/InitAllExtensions.h
    mlir/include/mlir/Target/LLVM/NVVM/Target.h
    mlir/include/mlir/Target/LLVM/ROCDL/Target.h
    mlir/include/mlir/Target/LLVMIR/Dialect/All.h
    mlir/include/mlir/Target/LLVMIR/Dialect/GPU/GPUToLLVMIRTranslation.h
    mlir/lib/CAPI/RegisterEverything/RegisterEverything.cpp
    mlir/lib/Dialect/GPU/Transforms/ModuleToBinary.cpp
    mlir/lib/Dialect/GPU/Transforms/NVVMAttachTarget.cpp
    mlir/lib/Dialect/GPU/Transforms/ROCDLAttachTarget.cpp
    mlir/lib/Dialect/GPU/Transforms/SerializeToBlob.cpp
    mlir/lib/Dialect/GPU/Transforms/SerializeToCubin.cpp
    mlir/lib/Dialect/GPU/Transforms/SerializeToHsaco.cpp
    mlir/lib/Target/LLVM/NVVM/Target.cpp
    mlir/lib/Target/LLVM/ROCDL/Target.cpp
    mlir/lib/Target/LLVMIR/ConvertToLLVMIR.cpp
    mlir/lib/Target/LLVMIR/Dialect/GPU/GPUToLLVMIRTranslation.cpp
    mlir/lib/Target/LLVMIR/Dialect/GPU/SelectObjectAttr.cpp
    mlir/tools/mlir-opt/mlir-opt.cpp
    mlir/unittests/Target/LLVM/SerializeNVVMTarget.cpp
    mlir/unittests/Target/LLVM/SerializeROCDLTarget.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/examples/transform/Ch2/transform-opt/transform-opt.cpp b/mlir/examples/transform/Ch2/transform-opt/transform-opt.cpp
index a875f630ef1b00..3a975313f93ff0 100644
--- a/mlir/examples/transform/Ch2/transform-opt/transform-opt.cpp
+++ b/mlir/examples/transform/Ch2/transform-opt/transform-opt.cpp
@@ -15,6 +15,7 @@
 #include "mlir/IR/DialectRegistry.h"
 #include "mlir/IR/MLIRContext.h"
 #include "mlir/InitAllDialects.h"
+#include "mlir/InitAllExtensions.h"
 #include "mlir/Tools/mlir-opt/MlirOptMain.h"
 #include "mlir/Transforms/Passes.h"
 #include <cstdlib>
@@ -35,6 +36,7 @@ int main(int argc, char **argv) {
   // Register all "core" dialects and our transform dialect extension.
   mlir::DialectRegistry registry;
   mlir::registerAllDialects(registry);
+  mlir::registerAllExtensions(registry);
   registerMyExtension(registry);
 
   // Register a handful of cleanup passes that we can run to make the output IR

diff  --git a/mlir/examples/transform/Ch3/transform-opt/transform-opt.cpp b/mlir/examples/transform/Ch3/transform-opt/transform-opt.cpp
index d0da0c803b77fd..1e4367ad469025 100644
--- a/mlir/examples/transform/Ch3/transform-opt/transform-opt.cpp
+++ b/mlir/examples/transform/Ch3/transform-opt/transform-opt.cpp
@@ -15,6 +15,7 @@
 #include "mlir/IR/DialectRegistry.h"
 #include "mlir/IR/MLIRContext.h"
 #include "mlir/InitAllDialects.h"
+#include "mlir/InitAllExtensions.h"
 #include "mlir/Tools/mlir-opt/MlirOptMain.h"
 #include "mlir/Transforms/Passes.h"
 #include <cstdlib>
@@ -35,6 +36,7 @@ int main(int argc, char **argv) {
   // Register all "core" dialects and our transform dialect extension.
   mlir::DialectRegistry registry;
   mlir::registerAllDialects(registry);
+  mlir::registerAllExtensions(registry);
   registerMyExtension(registry);
 
   // Register a handful of cleanup passes that we can run to make the output IR

diff  --git a/mlir/include/mlir/Dialect/GPU/Transforms/Passes.h b/mlir/include/mlir/Dialect/GPU/Transforms/Passes.h
index f0675dade3a8b6..033e8755501f96 100644
--- a/mlir/include/mlir/Dialect/GPU/Transforms/Passes.h
+++ b/mlir/include/mlir/Dialect/GPU/Transforms/Passes.h
@@ -87,8 +87,6 @@ class SerializeToBlobPass : public OperationPass<gpu::GPUModuleOp> {
   void runOnOperation() final;
 
 protected:
-  void getDependentDialects(DialectRegistry &registry) const override;
-
   /// Hook allowing the application of optimizations before codegen
   /// By default, does nothing
   virtual LogicalResult optimizeLlvm(llvm::Module &llvmModule,

diff  --git a/mlir/include/mlir/InitAllDialects.h b/mlir/include/mlir/InitAllDialects.h
index af67d58e795ba7..aa0a580aceb828 100644
--- a/mlir/include/mlir/InitAllDialects.h
+++ b/mlir/include/mlir/InitAllDialects.h
@@ -18,7 +18,6 @@
 #include "mlir/Dialect/AMX/AMXDialect.h"
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Affine/IR/ValueBoundsOpInterfaceImpl.h"
-#include "mlir/Dialect/Affine/TransformOps/AffineTransformOps.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Arith/IR/ValueBoundsOpInterfaceImpl.h"
 #include "mlir/Dialect/Arith/Transforms/BufferizableOpInterfaceImpl.h"
@@ -27,16 +26,13 @@
 #include "mlir/Dialect/ArmSVE/ArmSVEDialect.h"
 #include "mlir/Dialect/Async/IR/Async.h"
 #include "mlir/Dialect/Bufferization/IR/Bufferization.h"
-#include "mlir/Dialect/Bufferization/TransformOps/BufferizationTransformOps.h"
 #include "mlir/Dialect/Bufferization/Transforms/FuncBufferizableOpInterfaceImpl.h"
 #include "mlir/Dialect/Complex/IR/Complex.h"
 #include "mlir/Dialect/ControlFlow/IR/ControlFlow.h"
 #include "mlir/Dialect/DLTI/DLTI.h"
 #include "mlir/Dialect/EmitC/IR/EmitC.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
-#include "mlir/Dialect/Func/TransformOps/FuncTransformOps.h"
 #include "mlir/Dialect/GPU/IR/GPUDialect.h"
-#include "mlir/Dialect/GPU/TransformOps/GPUTransformOps.h"
 #include "mlir/Dialect/IRDL/IR/IRDL.h"
 #include "mlir/Dialect/Index/IR/IndexDialect.h"
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
@@ -44,7 +40,6 @@
 #include "mlir/Dialect/LLVMIR/ROCDLDialect.h"
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/Linalg/IR/ValueBoundsOpInterfaceImpl.h"
-#include "mlir/Dialect/Linalg/TransformOps/DialectExtension.h"
 #include "mlir/Dialect/Linalg/Transforms/BufferizableOpInterfaceImpl.h"
 #include "mlir/Dialect/Linalg/Transforms/TilingInterfaceImpl.h"
 #include "mlir/Dialect/MLProgram/IR/MLProgram.h"
@@ -52,11 +47,9 @@
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/MemRef/IR/MemRefMemorySlot.h"
 #include "mlir/Dialect/MemRef/IR/ValueBoundsOpInterfaceImpl.h"
-#include "mlir/Dialect/MemRef/TransformOps/MemRefTransformOps.h"
 #include "mlir/Dialect/MemRef/Transforms/BufferizableOpInterfaceImpl.h"
 #include "mlir/Dialect/MemRef/Transforms/RuntimeOpVerification.h"
 #include "mlir/Dialect/NVGPU/IR/NVGPUDialect.h"
-#include "mlir/Dialect/NVGPU/TransformOps/NVGPUTransformOps.h"
 #include "mlir/Dialect/OpenACC/OpenACC.h"
 #include "mlir/Dialect/OpenMP/OpenMPDialect.h"
 #include "mlir/Dialect/PDL/IR/PDL.h"
@@ -64,13 +57,11 @@
 #include "mlir/Dialect/Quant/QuantOps.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/Dialect/SCF/IR/ValueBoundsOpInterfaceImpl.h"
-#include "mlir/Dialect/SCF/TransformOps/SCFTransformOps.h"
 #include "mlir/Dialect/SCF/Transforms/BufferizableOpInterfaceImpl.h"
 #include "mlir/Dialect/SPIRV/IR/SPIRVDialect.h"
 #include "mlir/Dialect/Shape/IR/Shape.h"
 #include "mlir/Dialect/Shape/Transforms/BufferizableOpInterfaceImpl.h"
 #include "mlir/Dialect/SparseTensor/IR/SparseTensor.h"
-#include "mlir/Dialect/SparseTensor/TransformOps/SparseTensorTransformOps.h"
 #include "mlir/Dialect/SparseTensor/Transforms/BufferizableOpInterfaceImpl.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Dialect/Tensor/IR/TensorInferTypeOpInterfaceImpl.h"
@@ -83,11 +74,12 @@
 #include "mlir/Dialect/Transform/PDLExtension/PDLExtension.h"
 #include "mlir/Dialect/UB/IR/UBOps.h"
 #include "mlir/Dialect/Vector/IR/VectorOps.h"
-#include "mlir/Dialect/Vector/TransformOps/VectorTransformOps.h"
 #include "mlir/Dialect/Vector/Transforms/BufferizableOpInterfaceImpl.h"
 #include "mlir/Dialect/X86Vector/X86VectorDialect.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/Interfaces/CastInterfaces.h"
+#include "mlir/Target/LLVM/NVVM/Target.h"
+#include "mlir/Target/LLVM/ROCDL/Target.h"
 
 namespace mlir {
 
@@ -136,20 +128,6 @@ inline void registerAllDialects(DialectRegistry &registry) {
                   x86vector::X86VectorDialect>();
   // clang-format on
 
-  // Register all dialect extensions.
-  affine::registerTransformDialectExtension(registry);
-  bufferization::registerTransformDialectExtension(registry);
-  func::registerTransformDialectExtension(registry);
-  gpu::registerTransformDialectExtension(registry);
-  linalg::registerTransformDialectExtension(registry);
-  memref::registerTransformDialectExtension(registry);
-  nvgpu::registerTransformDialectExtension(registry);
-  scf::registerTransformDialectExtension(registry);
-  sparse_tensor::registerTransformDialectExtension(registry);
-  tensor::registerTransformDialectExtension(registry);
-  transform::registerPDLExtension(registry);
-  vector::registerTransformDialectExtension(registry);
-
   // Register all external models.
   affine::registerValueBoundsOpInterfaceExternalModels(registry);
   arith::registerBufferizableOpInterfaceExternalModels(registry);
@@ -174,6 +152,8 @@ inline void registerAllDialects(DialectRegistry &registry) {
   tensor::registerTilingInterfaceExternalModels(registry);
   tensor::registerValueBoundsOpInterfaceExternalModels(registry);
   vector::registerBufferizableOpInterfaceExternalModels(registry);
+  NVVM::registerNVVMTargetInterfaceExternalModels(registry);
+  ROCDL::registerROCDLTargetInterfaceExternalModels(registry);
 }
 
 /// Append all the MLIR dialects to the registry contained in the given context.

diff  --git a/mlir/include/mlir/InitAllExtensions.h b/mlir/include/mlir/InitAllExtensions.h
index fed2e15c4efb92..36899815866f62 100644
--- a/mlir/include/mlir/InitAllExtensions.h
+++ b/mlir/include/mlir/InitAllExtensions.h
@@ -23,9 +23,22 @@
 #include "mlir/Conversion/MemRefToLLVM/MemRefToLLVM.h"
 #include "mlir/Conversion/NVVMToLLVM/NVVMToLLVM.h"
 #include "mlir/Conversion/UBToLLVM/UBToLLVM.h"
+#include "mlir/Dialect/Affine/TransformOps/AffineTransformOps.h"
+#include "mlir/Dialect/Bufferization/TransformOps/BufferizationTransformOps.h"
 #include "mlir/Dialect/Func/Extensions/AllExtensions.h"
-#include "mlir/Target/LLVM/NVVM/Target.h"
-#include "mlir/Target/LLVM/ROCDL/Target.h"
+#include "mlir/Dialect/Func/TransformOps/FuncTransformOps.h"
+#include "mlir/Dialect/GPU/TransformOps/GPUTransformOps.h"
+#include "mlir/Dialect/Linalg/TransformOps/DialectExtension.h"
+#include "mlir/Dialect/MemRef/TransformOps/MemRefTransformOps.h"
+#include "mlir/Dialect/NVGPU/TransformOps/NVGPUTransformOps.h"
+#include "mlir/Dialect/SCF/TransformOps/SCFTransformOps.h"
+#include "mlir/Dialect/SparseTensor/TransformOps/SparseTensorTransformOps.h"
+#include "mlir/Dialect/Tensor/TransformOps/TensorTransformOps.h"
+#include "mlir/Dialect/Vector/TransformOps/VectorTransformOps.h"
+#include "mlir/Target/LLVMIR/Dialect/Builtin/BuiltinToLLVMIRTranslation.h"
+#include "mlir/Target/LLVMIR/Dialect/GPU/GPUToLLVMIRTranslation.h"
+#include "mlir/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.h"
+#include "mlir/Target/LLVMIR/Dialect/ROCDL/ROCDLToLLVMIRTranslation.h"
 
 #include <cstdlib>
 
@@ -37,6 +50,7 @@ namespace mlir {
 /// individually register the specific extensions that are useful for the
 /// pipelines and transformations you are using.
 inline void registerAllExtensions(DialectRegistry &registry) {
+  // Register all conversions to LLVM extensions.
   arith::registerConvertArithToLLVMInterface(registry);
   registerConvertComplexToLLVMInterface(registry);
   cf::registerConvertControlFlowToLLVMInterface(registry);
@@ -47,8 +61,23 @@ inline void registerAllExtensions(DialectRegistry &registry) {
   registerConvertMemRefToLLVMInterface(registry);
   registerConvertNVVMToLLVMInterface(registry);
   ub::registerConvertUBToLLVMInterface(registry);
-  registerNVVMTarget(registry);
-  registerROCDLTarget(registry);
+
+  // Register all transform dialect extensions.
+  affine::registerTransformDialectExtension(registry);
+  bufferization::registerTransformDialectExtension(registry);
+  func::registerTransformDialectExtension(registry);
+  gpu::registerTransformDialectExtension(registry);
+  linalg::registerTransformDialectExtension(registry);
+  memref::registerTransformDialectExtension(registry);
+  nvgpu::registerTransformDialectExtension(registry);
+  scf::registerTransformDialectExtension(registry);
+  sparse_tensor::registerTransformDialectExtension(registry);
+  tensor::registerTransformDialectExtension(registry);
+  transform::registerPDLExtension(registry);
+  vector::registerTransformDialectExtension(registry);
+
+  // Translation extensions need to be registered by calling
+  // `registerAllToLLVMIRTranslations` (see All.h).
 }
 
 } // namespace mlir

diff  --git a/mlir/include/mlir/Target/LLVM/NVVM/Target.h b/mlir/include/mlir/Target/LLVM/NVVM/Target.h
index ab011d34d6cb22..e3cb758725c8f2 100644
--- a/mlir/include/mlir/Target/LLVM/NVVM/Target.h
+++ b/mlir/include/mlir/Target/LLVM/NVVM/Target.h
@@ -16,13 +16,15 @@
 namespace mlir {
 class DialectRegistry;
 class MLIRContext;
+namespace NVVM {
 /// Registers the `TargetAttrInterface` for the `#nvvm.target` attribute in the
 /// given registry.
-void registerNVVMTarget(DialectRegistry &registry);
+void registerNVVMTargetInterfaceExternalModels(DialectRegistry &registry);
 
 /// Registers the `TargetAttrInterface` for the `#nvvm.target` attribute in the
 /// registry associated with the given context.
-void registerNVVMTarget(MLIRContext &context);
+void registerNVVMTargetInterfaceExternalModels(MLIRContext &context);
+} // namespace NVVM
 } // namespace mlir
 
 #endif // MLIR_TARGET_LLVM_NVVM_TARGET_H

diff  --git a/mlir/include/mlir/Target/LLVM/ROCDL/Target.h b/mlir/include/mlir/Target/LLVM/ROCDL/Target.h
index f6a548e550868c..4e39848bd42342 100644
--- a/mlir/include/mlir/Target/LLVM/ROCDL/Target.h
+++ b/mlir/include/mlir/Target/LLVM/ROCDL/Target.h
@@ -16,13 +16,15 @@
 namespace mlir {
 class DialectRegistry;
 class MLIRContext;
+namespace ROCDL {
 /// Registers the `TargetAttrInterface` for the `#rocdl.target` attribute in the
 /// given registry.
-void registerROCDLTarget(DialectRegistry &registry);
+void registerROCDLTargetInterfaceExternalModels(DialectRegistry &registry);
 
 /// Registers the `TargetAttrInterface` for the `#rocdl.target` attribute in the
 /// registry associated with the given context.
-void registerROCDLTarget(MLIRContext &context);
+void registerROCDLTargetInterfaceExternalModels(MLIRContext &context);
+} // namespace ROCDL
 } // namespace mlir
 
 #endif // MLIR_TARGET_LLVM_ROCDL_TARGET_H

diff  --git a/mlir/include/mlir/Target/LLVMIR/Dialect/All.h b/mlir/include/mlir/Target/LLVMIR/Dialect/All.h
index 65c1c515d4443b..895a37a63e9ff3 100644
--- a/mlir/include/mlir/Target/LLVMIR/Dialect/All.h
+++ b/mlir/include/mlir/Target/LLVMIR/Dialect/All.h
@@ -14,6 +14,8 @@
 #ifndef MLIR_TARGET_LLVMIR_DIALECT_ALL_H
 #define MLIR_TARGET_LLVMIR_DIALECT_ALL_H
 
+#include "mlir/Target/LLVM/NVVM/Target.h"
+#include "mlir/Target/LLVM/ROCDL/Target.h"
 #include "mlir/Target/LLVMIR/Dialect/AMX/AMXToLLVMIRTranslation.h"
 #include "mlir/Target/LLVMIR/Dialect/ArmNeon/ArmNeonToLLVMIRTranslation.h"
 #include "mlir/Target/LLVMIR/Dialect/ArmSME/ArmSMEToLLVMIRTranslation.h"
@@ -46,6 +48,16 @@ static inline void registerAllToLLVMIRTranslations(DialectRegistry &registry) {
   registerOpenMPDialectTranslation(registry);
   registerROCDLDialectTranslation(registry);
   registerX86VectorDialectTranslation(registry);
+
+  // Extension required for translating GPU offloading Ops.
+  gpu::registerOffloadingLLVMTranslationInterfaceExternalModels(registry);
+
+  // GPU target attribute interfaces are not used during translation, however
+  // the IR fails to verify if they are not registered due to the promise
+  // mechanism.
+  // TODO: remove these.
+  NVVM::registerNVVMTargetInterfaceExternalModels(registry);
+  ROCDL::registerROCDLTargetInterfaceExternalModels(registry);
 }
 
 /// Registers all dialects that can be translated from LLVM IR and the

diff  --git a/mlir/include/mlir/Target/LLVMIR/Dialect/GPU/GPUToLLVMIRTranslation.h b/mlir/include/mlir/Target/LLVMIR/Dialect/GPU/GPUToLLVMIRTranslation.h
index b42158d9b1e57c..19a34c6c00742a 100644
--- a/mlir/include/mlir/Target/LLVMIR/Dialect/GPU/GPUToLLVMIRTranslation.h
+++ b/mlir/include/mlir/Target/LLVMIR/Dialect/GPU/GPUToLLVMIRTranslation.h
@@ -29,7 +29,7 @@ void registerGPUDialectTranslation(MLIRContext &context);
 namespace gpu {
 /// Registers the offloading LLVM translation interfaces for
 /// `gpu.select_object`.
-void registerOffloadingLLVMTranslationInterfacesExternalModels(
+void registerOffloadingLLVMTranslationInterfaceExternalModels(
     mlir::DialectRegistry &registry);
 } // namespace gpu
 

diff  --git a/mlir/lib/CAPI/RegisterEverything/RegisterEverything.cpp b/mlir/lib/CAPI/RegisterEverything/RegisterEverything.cpp
index b63899bd53a730..c1c4a418b25526 100644
--- a/mlir/lib/CAPI/RegisterEverything/RegisterEverything.cpp
+++ b/mlir/lib/CAPI/RegisterEverything/RegisterEverything.cpp
@@ -9,9 +9,11 @@
 #include "mlir-c/RegisterEverything.h"
 
 #include "mlir/CAPI/IR.h"
+#include "mlir/IR/MLIRContext.h"
 #include "mlir/InitAllDialects.h"
 #include "mlir/InitAllExtensions.h"
 #include "mlir/InitAllPasses.h"
+#include "mlir/Target/LLVMIR/Dialect/All.h"
 #include "mlir/Target/LLVMIR/Dialect/Builtin/BuiltinToLLVMIRTranslation.h"
 #include "mlir/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.h"
 
@@ -22,8 +24,9 @@ void mlirRegisterAllDialects(MlirDialectRegistry registry) {
 
 void mlirRegisterAllLLVMTranslations(MlirContext context) {
   auto &ctx = *unwrap(context);
-  mlir::registerBuiltinDialectTranslation(ctx);
-  mlir::registerLLVMDialectTranslation(ctx);
+  mlir::DialectRegistry registry;
+  mlir::registerAllToLLVMIRTranslations(registry);
+  ctx.appendDialectRegistry(registry);
 }
 
 void mlirRegisterAllPasses() { mlir::registerAllPasses(); }

diff  --git a/mlir/lib/Dialect/GPU/Transforms/ModuleToBinary.cpp b/mlir/lib/Dialect/GPU/Transforms/ModuleToBinary.cpp
index 4888c8f79cc1f6..ad21dcc307fa24 100644
--- a/mlir/lib/Dialect/GPU/Transforms/ModuleToBinary.cpp
+++ b/mlir/lib/Dialect/GPU/Transforms/ModuleToBinary.cpp
@@ -15,11 +15,10 @@
 
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/GPU/IR/GPUDialect.h"
+#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
+#include "mlir/Dialect/LLVMIR/NVVMDialect.h"
+#include "mlir/Dialect/LLVMIR/ROCDLDialect.h"
 #include "mlir/IR/BuiltinOps.h"
-#include "mlir/Target/LLVM/NVVM/Target.h"
-#include "mlir/Target/LLVM/ROCDL/Target.h"
-#include "mlir/Target/LLVMIR/Dialect/GPU/GPUToLLVMIRTranslation.h"
-#include "mlir/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 
 #include "llvm/ADT/STLExtras.h"
@@ -46,13 +45,13 @@ class GpuModuleToBinaryPass
 void GpuModuleToBinaryPass::getDependentDialects(
     DialectRegistry &registry) const {
   // Register all GPU related translations.
-  registerLLVMDialectTranslation(registry);
-  registerGPUDialectTranslation(registry);
+  registry.insert<gpu::GPUDialect>();
+  registry.insert<LLVM::LLVMDialect>();
 #if MLIR_CUDA_CONVERSIONS_ENABLED == 1
-  registerNVVMTarget(registry);
+  registry.insert<NVVM::NVVMDialect>();
 #endif
 #if MLIR_ROCM_CONVERSIONS_ENABLED == 1
-  registerROCDLTarget(registry);
+  registry.insert<ROCDL::ROCDLDialect>();
 #endif
 }
 

diff  --git a/mlir/lib/Dialect/GPU/Transforms/NVVMAttachTarget.cpp b/mlir/lib/Dialect/GPU/Transforms/NVVMAttachTarget.cpp
index 1167002d9282ff..f159cdedfaea10 100644
--- a/mlir/lib/Dialect/GPU/Transforms/NVVMAttachTarget.cpp
+++ b/mlir/lib/Dialect/GPU/Transforms/NVVMAttachTarget.cpp
@@ -38,7 +38,7 @@ struct NVVMAttachTarget
   void runOnOperation() override;
 
   void getDependentDialects(DialectRegistry &registry) const override {
-    registerNVVMTarget(registry);
+    registry.insert<NVVM::NVVMDialect>();
   }
 };
 } // namespace

diff  --git a/mlir/lib/Dialect/GPU/Transforms/ROCDLAttachTarget.cpp b/mlir/lib/Dialect/GPU/Transforms/ROCDLAttachTarget.cpp
index 934ee666bef5e7..4f500fd337f4ce 100644
--- a/mlir/lib/Dialect/GPU/Transforms/ROCDLAttachTarget.cpp
+++ b/mlir/lib/Dialect/GPU/Transforms/ROCDLAttachTarget.cpp
@@ -38,7 +38,7 @@ struct ROCDLAttachTarget
   void runOnOperation() override;
 
   void getDependentDialects(DialectRegistry &registry) const override {
-    registerROCDLTarget(registry);
+    registry.insert<ROCDL::ROCDLDialect>();
   }
 };
 } // namespace

diff  --git a/mlir/lib/Dialect/GPU/Transforms/SerializeToBlob.cpp b/mlir/lib/Dialect/GPU/Transforms/SerializeToBlob.cpp
index 97aba904436752..2758689045df21 100644
--- a/mlir/lib/Dialect/GPU/Transforms/SerializeToBlob.cpp
+++ b/mlir/lib/Dialect/GPU/Transforms/SerializeToBlob.cpp
@@ -12,7 +12,9 @@
 //
 //===----------------------------------------------------------------------===//
 
+#include "mlir/Dialect/GPU/IR/GPUDialect.h"
 #include "mlir/Dialect/GPU/Transforms/Passes.h"
+#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "mlir/ExecutionEngine/OptUtils.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Target/LLVMIR/Dialect/GPU/GPUToLLVMIRTranslation.h"
@@ -23,8 +25,8 @@
 #include "llvm/Support/TargetSelect.h"
 #include "llvm/Target/TargetMachine.h"
 
-#include <string>
 #include <optional>
+#include <string>
 
 #define DEBUG_TYPE "serialize-to-blob"
 
@@ -124,13 +126,6 @@ gpu::SerializeToBlobPass::optimizeLlvm(llvm::Module &llvmModule,
   return success();
 }
 
-void gpu::SerializeToBlobPass::getDependentDialects(
-    DialectRegistry &registry) const {
-  registerGPUDialectTranslation(registry);
-  registerLLVMDialectTranslation(registry);
-  OperationPass<gpu::GPUModuleOp>::getDependentDialects(registry);
-}
-
 std::unique_ptr<llvm::TargetMachine>
 gpu::SerializeToBlobPass::createTargetMachine() {
   Location loc = getOperation().getLoc();

diff  --git a/mlir/lib/Dialect/GPU/Transforms/SerializeToCubin.cpp b/mlir/lib/Dialect/GPU/Transforms/SerializeToCubin.cpp
index f1c67019b262ff..6ea7f57ad3c227 100644
--- a/mlir/lib/Dialect/GPU/Transforms/SerializeToCubin.cpp
+++ b/mlir/lib/Dialect/GPU/Transforms/SerializeToCubin.cpp
@@ -12,6 +12,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/GPU/Transforms/Passes.h"
+#include "mlir/Dialect/LLVMIR/NVVMDialect.h"
 #include "llvm/Support/Debug.h"
 
 #if MLIR_GPU_TO_CUBIN_PASS_ENABLE
@@ -63,8 +64,6 @@ class SerializeToCubinPass
   }
 
 private:
-  void getDependentDialects(DialectRegistry &registry) const override;
-
   // Serializes PTX to CUBIN.
   std::unique_ptr<std::vector<char>>
   serializeISA(const std::string &isa) override;
@@ -100,12 +99,6 @@ SerializeToCubinPass::SerializeToCubinPass(StringRef triple, StringRef chip,
     this->optLevel.setValue(optLevel);
 }
 
-void SerializeToCubinPass::getDependentDialects(
-    DialectRegistry &registry) const {
-  registerNVVMDialectTranslation(registry);
-  gpu::SerializeToBlobPass::getDependentDialects(registry);
-}
-
 std::unique_ptr<std::vector<char>>
 SerializeToCubinPass::serializeISA(const std::string &isa) {
   Location loc = getOperation().getLoc();

diff  --git a/mlir/lib/Dialect/GPU/Transforms/SerializeToHsaco.cpp b/mlir/lib/Dialect/GPU/Transforms/SerializeToHsaco.cpp
index e204a232ed0abe..9db0a942f7c38f 100644
--- a/mlir/lib/Dialect/GPU/Transforms/SerializeToHsaco.cpp
+++ b/mlir/lib/Dialect/GPU/Transforms/SerializeToHsaco.cpp
@@ -12,6 +12,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/GPU/Transforms/Passes.h"
+#include "mlir/Dialect/LLVMIR/ROCDLDialect.h"
 #include "mlir/IR/Location.h"
 #include "mlir/IR/MLIRContext.h"
 
@@ -84,8 +85,6 @@ class SerializeToHsacoPass
   translateToLLVMIR(llvm::LLVMContext &llvmContext) override;
 
 private:
-  void getDependentDialects(DialectRegistry &registry) const override;
-
   // Loads LLVM bitcode libraries
   std::optional<SmallVector<std::unique_ptr<llvm::Module>, 3>>
   loadLibraries(SmallVectorImpl<char> &path,
@@ -145,12 +144,6 @@ SerializeToHsacoPass::SerializeToHsacoPass(StringRef triple, StringRef arch,
     this->optLevel.setValue(optLevel);
 }
 
-void SerializeToHsacoPass::getDependentDialects(
-    DialectRegistry &registry) const {
-  registerROCDLDialectTranslation(registry);
-  gpu::SerializeToBlobPass::getDependentDialects(registry);
-}
-
 std::optional<SmallVector<std::unique_ptr<llvm::Module>, 3>>
 SerializeToHsacoPass::loadLibraries(SmallVectorImpl<char> &path,
                                     SmallVectorImpl<StringRef> &libraries,

diff  --git a/mlir/lib/Target/LLVM/NVVM/Target.cpp b/mlir/lib/Target/LLVM/NVVM/Target.cpp
index 7c85ae3a12c637..77352cb3bc6c67 100644
--- a/mlir/lib/Target/LLVM/NVVM/Target.cpp
+++ b/mlir/lib/Target/LLVM/NVVM/Target.cpp
@@ -51,16 +51,17 @@ class NVVMTargetAttrImpl
 } // namespace
 
 // Register the NVVM dialect, the NVVM translation & the target interface.
-void mlir::registerNVVMTarget(DialectRegistry &registry) {
-  registerNVVMDialectTranslation(registry);
+void mlir::NVVM::registerNVVMTargetInterfaceExternalModels(
+    DialectRegistry &registry) {
   registry.addExtension(+[](MLIRContext *ctx, NVVM::NVVMDialect *dialect) {
     NVVMTargetAttr::attachInterface<NVVMTargetAttrImpl>(*ctx);
   });
 }
 
-void mlir::registerNVVMTarget(MLIRContext &context) {
+void mlir::NVVM::registerNVVMTargetInterfaceExternalModels(
+    MLIRContext &context) {
   DialectRegistry registry;
-  registerNVVMTarget(registry);
+  registerNVVMTargetInterfaceExternalModels(registry);
   context.appendDialectRegistry(registry);
 }
 

diff  --git a/mlir/lib/Target/LLVM/ROCDL/Target.cpp b/mlir/lib/Target/LLVM/ROCDL/Target.cpp
index 0c08fdc719611f..fe48f04fad434b 100644
--- a/mlir/lib/Target/LLVM/ROCDL/Target.cpp
+++ b/mlir/lib/Target/LLVM/ROCDL/Target.cpp
@@ -64,16 +64,17 @@ class ROCDLTargetAttrImpl
 } // namespace
 
 // Register the ROCDL dialect, the ROCDL translation and the target interface.
-void mlir::registerROCDLTarget(DialectRegistry &registry) {
-  registerROCDLDialectTranslation(registry);
+void mlir::ROCDL::registerROCDLTargetInterfaceExternalModels(
+    DialectRegistry &registry) {
   registry.addExtension(+[](MLIRContext *ctx, ROCDL::ROCDLDialect *dialect) {
     ROCDLTargetAttr::attachInterface<ROCDLTargetAttrImpl>(*ctx);
   });
 }
 
-void mlir::registerROCDLTarget(MLIRContext &context) {
+void mlir::ROCDL::registerROCDLTargetInterfaceExternalModels(
+    MLIRContext &context) {
   DialectRegistry registry;
-  registerROCDLTarget(registry);
+  registerROCDLTargetInterfaceExternalModels(registry);
   context.appendDialectRegistry(registry);
 }
 

diff  --git a/mlir/lib/Target/LLVMIR/ConvertToLLVMIR.cpp b/mlir/lib/Target/LLVMIR/ConvertToLLVMIR.cpp
index b7c1c40e13126a..45588937795348 100644
--- a/mlir/lib/Target/LLVMIR/ConvertToLLVMIR.cpp
+++ b/mlir/lib/Target/LLVMIR/ConvertToLLVMIR.cpp
@@ -13,8 +13,6 @@
 #include "mlir/Dialect/DLTI/DLTI.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/IR/BuiltinOps.h"
-#include "mlir/Target/LLVM/NVVM/Target.h"
-#include "mlir/Target/LLVM/ROCDL/Target.h"
 #include "mlir/Target/LLVMIR/Dialect/All.h"
 #include "mlir/Target/LLVMIR/Export.h"
 #include "mlir/Tools/mlir-translate/Translation.h"
@@ -38,8 +36,6 @@ void registerToLLVMIRTranslation() {
       },
       [](DialectRegistry &registry) {
         registry.insert<DLTIDialect, func::FuncDialect>();
-        registerNVVMTarget(registry);
-        registerROCDLTarget(registry);
         registerAllToLLVMIRTranslations(registry);
       });
 }

diff  --git a/mlir/lib/Target/LLVMIR/Dialect/GPU/GPUToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/GPU/GPUToLLVMIRTranslation.cpp
index ef98f737f07315..eecc8f1001ca4f 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/GPU/GPUToLLVMIRTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/GPU/GPUToLLVMIRTranslation.cpp
@@ -68,7 +68,6 @@ void mlir::registerGPUDialectTranslation(DialectRegistry &registry) {
   registry.addExtension(+[](MLIRContext *ctx, gpu::GPUDialect *dialect) {
     dialect->addInterfaces<GPUDialectLLVMIRTranslationInterface>();
   });
-  gpu::registerOffloadingLLVMTranslationInterfacesExternalModels(registry);
 }
 
 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 dd19fafa6ecb46..3b060ac1779db2 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/GPU/SelectObjectAttr.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/GPU/SelectObjectAttr.cpp
@@ -51,7 +51,7 @@ std::string getBinaryIdentifier(StringRef binaryName) {
 }
 } // namespace
 
-void mlir::gpu::registerOffloadingLLVMTranslationInterfacesExternalModels(
+void mlir::gpu::registerOffloadingLLVMTranslationInterfaceExternalModels(
     DialectRegistry &registry) {
   registry.addExtension(+[](MLIRContext *ctx, gpu::GPUDialect *dialect) {
     SelectObjectAttr::attachInterface<SelectObjectAttrImpl>(*ctx);

diff  --git a/mlir/tools/mlir-opt/mlir-opt.cpp b/mlir/tools/mlir-opt/mlir-opt.cpp
index 78bd70b40c91e7..b47634476c621b 100644
--- a/mlir/tools/mlir-opt/mlir-opt.cpp
+++ b/mlir/tools/mlir-opt/mlir-opt.cpp
@@ -19,6 +19,7 @@
 #include "mlir/Pass/Pass.h"
 #include "mlir/Pass/PassManager.h"
 #include "mlir/Support/FileUtilities.h"
+#include "mlir/Target/LLVMIR/Dialect/All.h"
 #include "mlir/Tools/mlir-opt/MlirOptMain.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/InitLLVM.h"
@@ -275,6 +276,7 @@ int main(int argc, char **argv) {
   DialectRegistry registry;
   registerAllDialects(registry);
   registerAllExtensions(registry);
+  registerAllToLLVMIRTranslations(registry);
 
 #ifdef MLIR_INCLUDE_TESTS
   ::test::registerTestDialect(registry);

diff  --git a/mlir/unittests/Target/LLVM/SerializeNVVMTarget.cpp b/mlir/unittests/Target/LLVM/SerializeNVVMTarget.cpp
index 7d95840fb9a52b..62c9b527e1e38c 100644
--- a/mlir/unittests/Target/LLVM/SerializeNVVMTarget.cpp
+++ b/mlir/unittests/Target/LLVM/SerializeNVVMTarget.cpp
@@ -15,6 +15,7 @@
 #include "mlir/Target/LLVMIR/Dialect/Builtin/BuiltinToLLVMIRTranslation.h"
 #include "mlir/Target/LLVMIR/Dialect/GPU/GPUToLLVMIRTranslation.h"
 #include "mlir/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.h"
+#include "mlir/Target/LLVMIR/Dialect/NVVM/NVVMToLLVMIRTranslation.h"
 
 #include "llvm/IRReader/IRReader.h"
 #include "llvm/Support/MemoryBufferRef.h"
@@ -40,7 +41,8 @@ class MLIRTargetLLVMNVVM : public ::testing::Test {
     registerBuiltinDialectTranslation(registry);
     registerLLVMDialectTranslation(registry);
     registerGPUDialectTranslation(registry);
-    registerNVVMTarget(registry);
+    registerNVVMDialectTranslation(registry);
+    NVVM::registerNVVMTargetInterfaceExternalModels(registry);
   }
 
   // Checks if PTXAS is in PATH.

diff  --git a/mlir/unittests/Target/LLVM/SerializeROCDLTarget.cpp b/mlir/unittests/Target/LLVM/SerializeROCDLTarget.cpp
index 9353c22bfa31f1..89cf5c5d2ada58 100644
--- a/mlir/unittests/Target/LLVM/SerializeROCDLTarget.cpp
+++ b/mlir/unittests/Target/LLVM/SerializeROCDLTarget.cpp
@@ -16,6 +16,7 @@
 #include "mlir/Target/LLVMIR/Dialect/Builtin/BuiltinToLLVMIRTranslation.h"
 #include "mlir/Target/LLVMIR/Dialect/GPU/GPUToLLVMIRTranslation.h"
 #include "mlir/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.h"
+#include "mlir/Target/LLVMIR/Dialect/ROCDL/ROCDLToLLVMIRTranslation.h"
 
 #include "llvm/IRReader/IRReader.h"
 #include "llvm/Support/FileSystem.h"
@@ -42,7 +43,8 @@ class MLIRTargetLLVMROCDL : public ::testing::Test {
     registerBuiltinDialectTranslation(registry);
     registerLLVMDialectTranslation(registry);
     registerGPUDialectTranslation(registry);
-    registerROCDLTarget(registry);
+    registerROCDLDialectTranslation(registry);
+    ROCDL::registerROCDLTargetInterfaceExternalModels(registry);
   }
 
   // Checks if a ROCm installation is available.

More information about the Mlir-commits mailing list