gpu-to-rocd filter-dialects

Author: Hardcode84

mlir/include/mlir/Conversion/Passes.td

@@ -578,20 +578,24 @@ def ConvertGpuOpsToROCDLOps : Pass<"convert-gpu-to-rocdl", "gpu::GPUModuleOp"> {
            /*default=*/"\"gfx000\"",
            "Chipset that these operations will run on">,
     Option<"indexBitwidth", "index-bitwidth", "unsigned",
-           /*default=kDeriveIndexBitwidthFromDataLayout*/"0",
+           /*default=kDeriveIndexBitwidthFromDataLayout*/ "0",
            "Bitwidth of the index type, 0 to use size of machine word">,
     Option<"useBarePtrCallConv", "use-bare-ptr-memref-call-conv", "bool",
            /*default=*/"false",
            "Replace memref arguments in GPU functions with bare pointers."
            "All memrefs must have static shape">,
     Option<"runtime", "runtime", "::mlir::gpu::amd::Runtime",
-          "::mlir::gpu::amd::Runtime::Unknown",
-          "Runtime code will be run on (default is Unknown, can also use HIP or OpenCl)",
-          [{::llvm::cl::values(
-            clEnumValN(::mlir::gpu::amd::Runtime::Unknown, "unknown", "Unknown (default)"),
-            clEnumValN(::mlir::gpu::amd::Runtime::HIP, "HIP", "HIP"),
-            clEnumValN(::mlir::gpu::amd::Runtime::OpenCL, "OpenCL", "OpenCL")
-          )}]>
+           "::mlir::gpu::amd::Runtime::Unknown",
+           "Runtime code will be run on (default is Unknown, can also use HIP "
+           "or OpenCl)",
+           [{::llvm::cl::values(
+               clEnumValN(::mlir::gpu::amd::Runtime::Unknown, "unknown",
+                          "Unknown (default)"),
+               clEnumValN(::mlir::gpu::amd::Runtime::HIP, "HIP", "HIP"),
+               clEnumValN(::mlir::gpu::amd::Runtime::OpenCL, "OpenCL",
+                          "OpenCL"))}]>,
+    ListOption<"filterDialects", "filter-dialects", "std::string",
+               "Run conversion patterns of only the specified dialects">,
   ];
 }
 

mlir/lib/Conversion/GPUToROCDL/LowerGpuOpsToROCDLOps.cpp

@@ -293,13 +293,33 @@ struct LowerGpuOpsToROCDLOpsPass
     RewritePatternSet llvmPatterns(ctx);
     LLVMConversionTarget target(getContext());
 
-    for (Dialect *dialect : ctx->getLoadedDialects()) {
-      auto iface = dyn_cast<ConvertToLLVMPatternInterface>(dialect);
-      if (!iface)
-        continue;
+    if (!filterDialects.empty()) {
+      for (StringRef dialectName : filterDialects) {
+        Dialect *dialect = ctx->getLoadedDialect(dialectName);
+        // Dialect may not be loaded if it wasn't used in source module, ignore.
+        if (!dialect)
+          continue;
+
+        auto *iface = dyn_cast<ConvertToLLVMPatternInterface>(dialect);
+        if (!iface) {
+          m.emitError()
+              << "dialect does not implement ConvertToLLVMPatternInterface: "
+              << dialectName << "\n";
+          return signalPassFailure();
+        }
 
-      iface->populateConvertToLLVMConversionPatterns(target, converter,
-                                                     llvmPatterns);
+        iface->populateConvertToLLVMConversionPatterns(target, converter,
+                                                       llvmPatterns);
+      }
+    } else {
+      for (Dialect *dialect : ctx->getLoadedDialects()) {
+        auto iface = dyn_cast<ConvertToLLVMPatternInterface>(dialect);
+        if (!iface)
+          continue;
+
+        iface->populateConvertToLLVMConversionPatterns(target, converter,
+                                                       llvmPatterns);
+      }
     }
 
     populateAMDGPUToROCDLConversionPatterns(converter, llvmPatterns,

mlir/test/Conversion/GPUToROCDL/gpu-to-rocdl.mlir

@@ -1,4 +1,5 @@
 // RUN: mlir-opt %s -convert-gpu-to-rocdl -split-input-file | FileCheck %s
+// RUN: mlir-opt %s -convert-gpu-to-rocdl='filter-dialects=func,arith,math' -split-input-file | FileCheck %s
 // RUN: mlir-opt %s -convert-gpu-to-rocdl='index-bitwidth=32' -split-input-file | FileCheck --check-prefix=CHECK32 %s
 
 // CHECK-LABEL: @test_module