Ensure IREE GPU dialect is registered for all GPU targets (fixes #17736…

…) (#17737) Before this commit, it was only registered for ROCm, so the compiler would sometimes crash during compilation (including for some things compiled as part of the IREE build process) when trying to use IREE GPU dialect attributes if the ROCm target wasn't enabled. Signed-off-by: Andrea Faulds <[email protected]>
iree-org · Jun 25, 2024 · 0d2c780 · 0d2c780
1 parent 1c53595
commit 0d2c780
Show file tree

Hide file tree

Showing 4 changed files with 6 additions and 4 deletions.
diff --git a/compiler/plugins/target/CUDA/CUDATarget.cpp b/compiler/plugins/target/CUDA/CUDATarget.cpp
@@ -450,7 +450,7 @@ class CUDATargetBackend final : public TargetBackend {
     // `LLVMGPULowerExecutableTargetPass`.
     registry.insert<gpu::GPUDialect, nvgpu::NVGPUDialect,
                     IREE::Codegen::IREECodegenDialect,
-                    transform::TransformDialect>();
+                    transform::TransformDialect, IREE::GPU::IREEGPUDialect>();
     mlir::registerBuiltinDialectTranslation(registry);
     mlir::registerLLVMDialectTranslation(registry);
     mlir::registerNVVMDialectTranslation(registry);

diff --git a/compiler/plugins/target/MetalSPIRV/MetalSPIRVTarget.cpp b/compiler/plugins/target/MetalSPIRV/MetalSPIRVTarget.cpp
@@ -112,7 +112,8 @@ class MetalSPIRVTargetBackend : public TargetBackend {
 
   void getDependentDialects(DialectRegistry &registry) const override {
     registry.insert<gpu::GPUDialect, IREE::Codegen::IREECodegenDialect,
-                    IREE::Flow::FlowDialect, spirv::SPIRVDialect>();
+                    IREE::Flow::FlowDialect, spirv::SPIRVDialect,
+                    IREE::GPU::IREEGPUDialect>();
   }
 
   void

diff --git a/compiler/plugins/target/VulkanSPIRV/VulkanSPIRVTarget.cpp b/compiler/plugins/target/VulkanSPIRV/VulkanSPIRVTarget.cpp
@@ -126,7 +126,7 @@ class VulkanSPIRVTargetBackend : public TargetBackend {
 
   void getDependentDialects(DialectRegistry &registry) const override {
     registry.insert<IREE::Codegen::IREECodegenDialect, spirv::SPIRVDialect,
-                    gpu::GPUDialect>();
+                    gpu::GPUDialect, IREE::GPU::IREEGPUDialect>();
   }
 
   void

diff --git a/compiler/plugins/target/WebGPUSPIRV/WebGPUSPIRVTarget.cpp b/compiler/plugins/target/WebGPUSPIRV/WebGPUSPIRVTarget.cpp
@@ -106,7 +106,8 @@ class WebGPUSPIRVTargetBackend : public TargetBackend {
   //     pipeline created by buildTranslationPassPipeline)
   void getDependentDialects(DialectRegistry &registry) const override {
     registry.insert<IREE::Codegen::IREECodegenDialect, IREE::Flow::FlowDialect,
-                    spirv::SPIRVDialect, gpu::GPUDialect>();
+                    spirv::SPIRVDialect, gpu::GPUDialect,
+                    IREE::GPU::IREEGPUDialect>();
   }
 
   void