[Mlir-commits] [mlir] [mlir][nvvm_target] Add nvvm options into NVVMAttachTarget pass (PR #126478)

[Zichen Lu]

https://github.com/MikaOvO updated https://github.com/llvm/llvm-project/pull/126478

>From 0e0db640b3a1ae01cc510f2f894b626e7c49bc6e Mon Sep 17 00:00:00 2001
From: Zichen Lu <mikaovo2000 at gmail.com>
Date: Mon, 10 Feb 2025 15:18:39 +0800
Subject: [PATCH] [mlir][nvvm_target] Add nvvm options into NVVMAttachTarget
 pass

---
 .../mlir/Dialect/GPU/IR/CompilationInterfaces.h   |  4 ++++
 .../include/mlir/Dialect/GPU/Transforms/Passes.td |  3 +++
 mlir/lib/Dialect/GPU/IR/GPUDialect.cpp            |  7 ++++++-
 .../Dialect/GPU/Transforms/NVVMAttachTarget.cpp   | 15 ++++++++++++++-
 mlir/test/Dialect/GPU/nvvm-attach-target.mlir     | 15 +++++++++++++++
 5 files changed, 42 insertions(+), 2 deletions(-)
 create mode 100644 mlir/test/Dialect/GPU/nvvm-attach-target.mlir

diff --git a/mlir/include/mlir/Dialect/GPU/IR/CompilationInterfaces.h b/mlir/include/mlir/Dialect/GPU/IR/CompilationInterfaces.h
index c950ef220f692f6..9a890ae24d8fc37 100644
--- a/mlir/include/mlir/Dialect/GPU/IR/CompilationInterfaces.h
+++ b/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>())`.
diff --git a/mlir/include/mlir/Dialect/GPU/Transforms/Passes.td b/mlir/include/mlir/Dialect/GPU/Transforms/Passes.td
index e055164a1c384ee..86459c506cf1055 100644
--- a/mlir/include/mlir/Dialect/GPU/Transforms/Passes.td
+++ b/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.">,
   ];
diff --git a/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp b/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp
index d06f10d3137a183..1bdeb3e356f4be9 100644
--- a/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp
+++ b/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"
diff --git a/mlir/lib/Dialect/GPU/Transforms/NVVMAttachTarget.cpp b/mlir/lib/Dialect/GPU/Transforms/NVVMAttachTarget.cpp
index dd705cd33831217..663bcbf7a848cea 100644
--- a/mlir/lib/Dialect/GPU/Transforms/NVVMAttachTarget.cpp
+++ b/mlir/lib/Dialect/GPU/Transforms/NVVMAttachTarget.cpp
@@ -45,7 +45,7 @@ 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));
   };
@@ -53,6 +53,19 @@ DictionaryAttr NVVMAttachTarget::getFlags(OpBuilder &builder) const {
     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;
diff --git a/mlir/test/Dialect/GPU/nvvm-attach-target.mlir b/mlir/test/Dialect/GPU/nvvm-attach-target.mlir
new file mode 100644
index 000000000000000..26b8885242b31fd
--- /dev/null
+++ b/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
+    }
+  }
+}
\ No newline at end of file