Adding the ELF section to the NVVMTargetAttr, to propagate to the gpu.binary Op

Author: Renaud-K

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

@@ -139,6 +139,9 @@ def GpuNVVMAttachTarget: Pass<"nvvm-attach-target", ""> {
     Option<"ftzFlag", "ftz", "bool",
            /*default=*/"false",
            "Enable flush to zero for denormals.">,
+    Option<"elfSection", "section", "std::string",
+           /*default=*/"\"\"",
+           "ELF section where the module needs to be created.">,
     ListOption<"linkLibs", "l", "std::string",
            "Extra bitcode libraries paths to link to.">,
   ];

mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td

@@ -2266,20 +2266,22 @@ def NVVM_TargettAttr : NVVM_Attr<"NVVMTarget", "target"> {
     StringRefParameter<"Target triple.", "\"nvptx64-nvidia-cuda\"">:$triple,
     StringRefParameter<"Target chip.", "\"sm_50\"">:$chip,
     StringRefParameter<"Target chip features.", "\"+ptx60\"">:$features,
+    OptionalParameter<"StringAttr", "ELF section.">:$section,
     OptionalParameter<"DictionaryAttr", "Target specific flags.">:$flags,
     OptionalParameter<"ArrayAttr", "Files to link to the LLVM module.">:$link
   );
   let assemblyFormat = [{
-    (`<` struct($O, $triple, $chip, $features, $flags, $link)^ `>`)?
+    (`<` struct($O, $triple, $chip, $features, $section, $flags, $link)^ `>`)?
   }];
   let builders = [
     AttrBuilder<(ins CArg<"int", "2">:$optLevel,
                      CArg<"StringRef", "\"nvptx64-nvidia-cuda\"">:$triple,
                      CArg<"StringRef", "\"sm_50\"">:$chip,
                      CArg<"StringRef", "\"+ptx60\"">:$features,
+                     CArg<"StringAttr", "nullptr">:$section,
                      CArg<"DictionaryAttr", "nullptr">:$targetFlags,
                      CArg<"ArrayAttr", "nullptr">:$linkFiles), [{
-      return Base::get($_ctxt, optLevel, triple, chip, features, targetFlags, linkFiles);
+      return Base::get($_ctxt, optLevel, triple, chip, features, section, targetFlags, linkFiles);
     }]>
   ];
   let skipDefaultBuilders = 1;

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

@@ -63,7 +63,9 @@ void NVVMAttachTarget::runOnOperation() {
   ArrayRef<std::string> libs(linkLibs);
   SmallVector<StringRef> filesToLink(libs);
   auto target = builder.getAttr<NVVMTargetAttr>(
-      optLevel, triple, chip, features, getFlags(builder),
+      optLevel, triple, chip, features,
+      elfSection.empty() ? nullptr : builder.getStringAttr(elfSection),
+      getFlags(builder),
       filesToLink.empty() ? nullptr : builder.getStrArrayAttr(filesToLink));
   llvm::Regex matcher(moduleMatcher);
   for (Region &region : getOperation()->getRegions())

mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp

@@ -1185,8 +1185,8 @@ LogicalResult NVVMDialect::verifyRegionArgAttribute(Operation *op,
 LogicalResult
 NVVMTargetAttr::verify(function_ref<InFlightDiagnostic()> emitError,
                        int optLevel, StringRef triple, StringRef chip,
-                       StringRef features, DictionaryAttr flags,
-                       ArrayAttr files) {
+                       StringRef features, StringAttr elfSection,
+                       DictionaryAttr flags, ArrayAttr files) {
   if (optLevel < 0 || optLevel > 3) {
     emitError() << "The optimization level must be a number between 0 and 3.";
     return failure();

mlir/lib/Target/LLVM/NVVM/Target.cpp

@@ -664,9 +664,17 @@ NVVMTargetAttrImpl::createObject(Attribute attribute, Operation *module,
   gpu::CompilationTarget format = options.getCompilationTarget();
   DictionaryAttr objectProps;
   Builder builder(attribute.getContext());
+  SmallVector<NamedAttribute, 2> properties;
   if (format == gpu::CompilationTarget::Assembly)
-    objectProps = builder.getDictionaryAttr(
-        {builder.getNamedAttr("O", builder.getI32IntegerAttr(target.getO()))});
+    properties.push_back(
+        builder.getNamedAttr("O", builder.getI32IntegerAttr(target.getO())));
+
+  if (StringAttr section = target.getSection())
+    properties.push_back(builder.getNamedAttr("section", section));
+
+  if (!properties.empty())
+    objectProps = builder.getDictionaryAttr(properties);
+
   return builder.getAttr<gpu::ObjectAttr>(
       attribute, format,
       builder.getStringAttr(StringRef(object.data(), object.size())),

mlir/test/Dialect/LLVMIR/attach-targets.mlir

@@ -1,13 +1,16 @@
 // RUN: mlir-opt %s --nvvm-attach-target='module=nvvm.* O=3 chip=sm_90' --rocdl-attach-target='module=rocdl.* O=3 chip=gfx90a' | FileCheck %s
 // RUN: mlir-opt %s --nvvm-attach-target='module=options.* O=1 chip=sm_70 fast=true ftz=true' --rocdl-attach-target='module=options.* l=file1.bc,file2.bc wave64=false finite-only=true' | FileCheck %s --check-prefix=CHECK_OPTS
+// RUN: mlir-opt %s --nvvm-attach-target='module=nvvm.* section=__fatbin' | FileCheck %s --check-prefix=CHECK_SECTION
 
 module attributes {gpu.container_module} {
 // Verify the target is appended.
 // CHECK: @nvvm_module_1 [#nvvm.target<O = 3, chip = "sm_90">] {
+// CHECK_SECTION: @nvvm_module_1 [#nvvm.target<section = "__fatbin">] 
 gpu.module @nvvm_module_1 {
 }
 // Verify the target is appended.
 // CHECK: @nvvm_module_2 [#nvvm.target<chip = "sm_60">, #nvvm.target<O = 3, chip = "sm_90">] {
+// CHECK_SECTION: gpu.module @nvvm_module_2 [#nvvm.target<chip = "sm_60">, #nvvm.target<section = "__fatbin">] 
 gpu.module @nvvm_module_2 [#nvvm.target<chip = "sm_60">] {
 }
 // Verify the target is not added multiple times.