[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

Author: nicolasvasilache

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

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

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,

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,50 +26,42 @@
 #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"
 #include "mlir/Dialect/LLVMIR/NVVMDialect.h"
 #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"
 #include "mlir/Dialect/Math/IR/Math.h"
 #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"
 #include "mlir/Dialect/PDLInterp/IR/PDLInterp.h"
 #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.

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

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

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

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

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
 

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(); }

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
 }
 

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

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

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();