[Mlir-commits] [mlir] Adding the ELF section to the NVVMTargetAttr, to propagate to the gpu.binary Op (PR #119440)

Renaud Kauffmann llvmlistbot at llvm.org
Tue Dec 10 11:32:29 PST 2024 

https://github.com/Renaud-K created https://github.com/llvm/llvm-project/pull/119440

This is a follow-up of #117246.

I thought then it would be easy to edit a DictionaryAttr but it turns out that these attributes are immutable and need to be passed during the construction of the gpu.binary Op. This requires the section to be part of the NVVMTargetAttr and to be an option to the GpuNVVMAttachTarget pass.



>From d62fa8dc4d60c960bdc58bffbc3c251573490703 Mon Sep 17 00:00:00 2001
From: Renaud-K <rkauffmann at nvidia.com>
Date: Tue, 10 Dec 2024 11:25:13 -0800
Subject: [PATCH] Adding the ELF section to the NVVMTargetAttr, to propagate to
 the gpu.binary Op

---
 mlir/include/mlir/Dialect/GPU/Transforms/Passes.td  |  3 +++
 mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td         |  6 ++++--
 .../lib/Dialect/GPU/Transforms/NVVMAttachTarget.cpp |  4 +++-
 mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp          |  2 +-
 mlir/lib/Target/LLVM/NVVM/Target.cpp                | 13 +++++++++++--
 mlir/test/Dialect/LLVMIR/attach-targets.mlir        |  3 +++
 6 files changed, 25 insertions(+), 6 deletions(-)

diff --git a/mlir/include/mlir/Dialect/GPU/Transforms/Passes.td b/mlir/include/mlir/Dialect/GPU/Transforms/Passes.td
index 4a9ddafdd177d2..784721d19d0ccd 100644
--- a/mlir/include/mlir/Dialect/GPU/Transforms/Passes.td
+++ b/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.">,
   ];
diff --git a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
index 296a3c305e5bf4..03c03d5f66b9af 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
+++ b/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;
diff --git a/mlir/lib/Dialect/GPU/Transforms/NVVMAttachTarget.cpp b/mlir/lib/Dialect/GPU/Transforms/NVVMAttachTarget.cpp
index dd705cd3383121..5628361c6a843c 100644
--- a/mlir/lib/Dialect/GPU/Transforms/NVVMAttachTarget.cpp
+++ b/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())
diff --git a/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp b/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
index ca04af0b060b4f..696cab52d41900 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
@@ -1185,7 +1185,7 @@ LogicalResult NVVMDialect::verifyRegionArgAttribute(Operation *op,
 LogicalResult
 NVVMTargetAttr::verify(function_ref<InFlightDiagnostic()> emitError,
                        int optLevel, StringRef triple, StringRef chip,
-                       StringRef features, DictionaryAttr flags,
+                       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.";
diff --git a/mlir/lib/Target/LLVM/NVVM/Target.cpp b/mlir/lib/Target/LLVM/NVVM/Target.cpp
index 3c92359915ded4..82ead374420d9e 100644
--- a/mlir/lib/Target/LLVM/NVVM/Target.cpp
+++ b/mlir/lib/Target/LLVM/NVVM/Target.cpp
@@ -664,9 +664,18 @@ 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())),
diff --git a/mlir/test/Dialect/LLVMIR/attach-targets.mlir b/mlir/test/Dialect/LLVMIR/attach-targets.mlir
index 83733db400798e..832e0443e720f7 100644
--- a/mlir/test/Dialect/LLVMIR/attach-targets.mlir
+++ b/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.

More information about the Mlir-commits mailing list