diff --git a/mlir/include/mlir/Dialect/GPU/IR/CompilationInterfaces.h b/mlir/include/mlir/Dialect/GPU/IR/CompilationInterfaces.h index c950ef220f692..9a890ae24d8fc 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> + tokenizeCmdOptions(const std::string &cmdOptions); + protected: /// Derived classes must use this constructor to initialize `typeID` to the /// appropiate value: ie. `TargetOptions(TypeID::get())`. diff --git a/mlir/include/mlir/Dialect/GPU/Transforms/Passes.td b/mlir/include/mlir/Dialect/GPU/Transforms/Passes.td index e055164a1c384..86459c506cf10 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 d06f10d3137a1..1bdeb3e356f4b 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> -TargetOptions::tokenizeCmdOptions() const { +TargetOptions::tokenizeCmdOptions(const std::string &cmdOptions) { std::pair> options; llvm::StringSaver stringSaver(options.first); StringRef opts = cmdOptions; @@ -2586,6 +2586,11 @@ TargetOptions::tokenizeCmdOptions() const { return options; } +std::pair> +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 dd705cd338312..663bcbf7a848c 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 flags; + SmallVector 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 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 0000000000000..8051ff3679b95 --- /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="nvvm-options=-opt=1" | FileCheck %s -check-prefix=CHECK-OPTIONS + +module attributes {gpu.container_module} { + // CHECK-LABEL:gpu.module @kernel_module1 + // CHECK: [#nvvm.target] + // CHECK-OPTIONS:[#nvvm.target] + 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