[mlir][gpu] GPUToROCDL: Add C++ argument to populate allowedDialects

The `convert-gpu-to-rocdl` pass provides the option `allowed-dialects`,
which allows users to control which dialects can be used to populate
conversions.

This PR adds a C++ argument to createLowerGpuOpsToROCDLOpsPass, so
that this option can also be controlled programatically when creating
the pass.

Author: pabloantoniom

mlir/include/mlir/Conversion/GPUToROCDL/GPUToROCDLPass.h

@@ -10,6 +10,8 @@
 
 #include "mlir/Conversion/GPUToROCDL/Runtimes.h"
 #include "mlir/Conversion/LLVMCommon/LoweringOptions.h"
+#include "llvm/ADT/DenseSet.h"
+#include <cstddef>
 #include <memory>
 
 namespace mlir {
@@ -50,7 +52,9 @@ createLowerGpuOpsToROCDLOpsPass(
     const std::string &chipset = "gfx900",
     unsigned indexBitwidth = kDeriveIndexBitwidthFromDataLayout,
     bool useBarePtrCallConv = false,
-    gpu::amd::Runtime runtime = gpu::amd::Runtime::Unknown);
+    gpu::amd::Runtime runtime = gpu::amd::Runtime::Unknown,
+    const std::optional<llvm::SmallDenseSet<llvm::StringRef>> &allowedDialects =
+        std::nullopt);
 
 } // namespace mlir
 

mlir/lib/Conversion/GPUToROCDL/LowerGpuOpsToROCDLOps.cpp

@@ -288,9 +288,10 @@ struct GPUShuffleOpLowering : public ConvertOpToLLVMPattern<gpu::ShuffleOp> {
 struct LowerGpuOpsToROCDLOpsPass final
     : public impl::ConvertGpuOpsToROCDLOpsBase<LowerGpuOpsToROCDLOpsPass> {
   LowerGpuOpsToROCDLOpsPass() = default;
-  LowerGpuOpsToROCDLOpsPass(const std::string &chipset, unsigned indexBitwidth,
-                            bool useBarePtrCallConv,
-                            gpu::amd::Runtime runtime) {
+  LowerGpuOpsToROCDLOpsPass(
+      const std::string &chipset, unsigned indexBitwidth,
+      bool useBarePtrCallConv, gpu::amd::Runtime runtime,
+      std::optional<llvm::SmallDenseSet<StringRef>> allowedDialects) {
     if (this->chipset.getNumOccurrences() == 0)
       this->chipset = chipset;
     if (this->indexBitwidth.getNumOccurrences() == 0)
@@ -299,6 +300,12 @@ struct LowerGpuOpsToROCDLOpsPass final
       this->useBarePtrCallConv = useBarePtrCallConv;
     if (this->runtime.getNumOccurrences() == 0)
       this->runtime = runtime;
+    if (this->allowedDialects.getNumOccurrences() == 0 &&
+        allowedDialects.has_value()) {
+      for (auto &str : allowedDialects.value()) {
+        this->allowedDialects.push_back(str.str());
+      }
+    }
   }
 
   void getDependentDialects(DialectRegistry &registry) const override {
@@ -501,10 +508,10 @@ void mlir::populateGpuToROCDLConversionPatterns(
 }
 
 std::unique_ptr<OperationPass<gpu::GPUModuleOp>>
-mlir::createLowerGpuOpsToROCDLOpsPass(const std::string &chipset,
-                                      unsigned indexBitwidth,
-                                      bool useBarePtrCallConv,
-                                      gpu::amd::Runtime runtime) {
+mlir::createLowerGpuOpsToROCDLOpsPass(
+    const std::string &chipset, unsigned indexBitwidth, bool useBarePtrCallConv,
+    gpu::amd::Runtime runtime,
+    const std::optional<llvm::SmallDenseSet<StringRef>> &allowedDialects) {
   return std::make_unique<LowerGpuOpsToROCDLOpsPass>(
-      chipset, indexBitwidth, useBarePtrCallConv, runtime);
+      chipset, indexBitwidth, useBarePtrCallConv, runtime, allowedDialects);
 }