[mlir][nvvm_target] Add nvvm options into NVVMAttachTarget pass

Author: MikaOvO

mlir/include/mlir/Dialect/GPU/IR/CompilationInterfaces.h

@@ -108,6 +108,10 @@ class TargetOptions {
   /// Returns the default compilation target: `CompilationTarget::Fatbin`.
   static CompilationTarget getDefaultCompilationTarget();
 
+  /// Returns a tokenization of the command line options.
+  static std::pair<llvm::BumpPtrAllocator, SmallVector<const char *>>
+  tokenizeCmdOptions(const std::string &cmdOptions);
+
 protected:
   /// Derived classes must use this constructor to initialize `typeID` to the
   /// appropiate value: ie. `TargetOptions(TypeID::get<DerivedClass>())`.

mlir/include/mlir/Dialect/GPU/Transforms/Passes.td

@@ -141,6 +141,9 @@ def GpuNVVMAttachTarget: Pass<"nvvm-attach-target", ""> {
     Option<"ftzFlag", "ftz", "bool",
            /*default=*/"false",
            "Enable flush to zero for denormals.">,
+    Option<"nvvmOptions", "nvvm-options", "std::string",
+           /*default=*/ [{""}],
+           "Options passed to libnvvm">,
     ListOption<"linkLibs", "l", "std::string",
            "Extra bitcode libraries paths to link to.">,
   ];

mlir/lib/Dialect/GPU/IR/GPUDialect.cpp

@@ -2564,7 +2564,7 @@ CompilationTarget TargetOptions::getDefaultCompilationTarget() {
 }
 
 std::pair<llvm::BumpPtrAllocator, SmallVector<const char *>>
-TargetOptions::tokenizeCmdOptions() const {
+TargetOptions::tokenizeCmdOptions(const std::string &cmdOptions) {
   std::pair<llvm::BumpPtrAllocator, SmallVector<const char *>> options;
   llvm::StringSaver stringSaver(options.first);
   StringRef opts = cmdOptions;
@@ -2586,6 +2586,11 @@ TargetOptions::tokenizeCmdOptions() const {
   return options;
 }
 
+std::pair<llvm::BumpPtrAllocator, SmallVector<const char *>>
+TargetOptions::tokenizeCmdOptions() const {
+  return tokenizeCmdOptions(cmdOptions);
+}
+
 MLIR_DEFINE_EXPLICIT_TYPE_ID(::mlir::gpu::TargetOptions)
 
 #include "mlir/Dialect/GPU/IR/GPUOpInterfaces.cpp.inc"

mlir/lib/Dialect/GPU/Transforms/NVVMAttachTarget.cpp

@@ -45,14 +45,27 @@ struct NVVMAttachTarget
 
 DictionaryAttr NVVMAttachTarget::getFlags(OpBuilder &builder) const {
   UnitAttr unitAttr = builder.getUnitAttr();
-  SmallVector<NamedAttribute, 2> flags;
+  SmallVector<NamedAttribute, 3> flags;
   auto addFlag = [&](StringRef flag) {
     flags.push_back(builder.getNamedAttr(flag, unitAttr));
   };
   if (fastFlag)
     addFlag("fast");
   if (ftzFlag)
     addFlag("ftz");
+  if (!nvvmOptions.empty()) {
+    auto options = gpu::TargetOptions::tokenizeCmdOptions(nvvmOptions);
+    if (!options.second.empty()) {
+      llvm::SmallVector<mlir::Attribute> nvvmOptionAttrs;
+      for (const char *opt : options.second) {
+        nvvmOptionAttrs.emplace_back(
+            mlir::StringAttr::get(builder.getContext(), StringRef(opt)));
+      }
+      flags.push_back(builder.getNamedAttr(
+          "libNVVMOptions",
+          mlir::ArrayAttr::get(builder.getContext(), nvvmOptionAttrs)));
+    }
+  }
   if (!flags.empty())
     return builder.getDictionaryAttr(flags);
   return nullptr;

mlir/test/Dialect/GPU/nvvm-attach-target.mlir

@@ -0,0 +1,15 @@
+// RUN: mlir-opt %s --nvvm-attach-target="" | FileCheck %s
+// RUN: mlir-opt %s --nvvm-attach-target="-opt=1" | FileCheck %s -check-prefix=CHECK-OPTIONS
+
+module attributes {gpu.container_module} {
+  // CHECK-LABEL:gpu.binary @kernel_module1
+  // CHECK:gpu.module @kernel_module1 [#nvvm.target]
+  // CHECK-OPTIONS:gpu.module @kernel_module1 [#nvvm.target<flags = {libNVVMOptions = ["-opt=1"]}>]
+  gpu.module @kernel_module1 {
+    llvm.func @kernel(%arg0: i32, %arg1: !llvm.ptr,
+        %arg2: !llvm.ptr, %arg3: i64, %arg4: i64,
+        %arg5: i64) attributes {gpu.kernel} {
+      llvm.return
+    }
+  }
+}