[Mlir-commits] [mlir] [mlir][gpu] Pass GPU module to `TargetAttrInterface::createObject`. (PR #94910)

[Fabian Mora]

https://github.com/fabianmcg updated https://github.com/llvm/llvm-project/pull/94910

>From 08914129ceccc27e3fc35c661761885c4f01d6cf Mon Sep 17 00:00:00 2001
From: Fabian Mora <fmora.dev at gmail.com>
Date: Sun, 9 Jun 2024 20:10:31 +0000
Subject: [PATCH 1/2] [mlir][gpu] Pass GPU module to
 `TargetAttrInterface::createObject`.

This patch adds an argument to `gpu::TargetAttrInterface::createObject` to pass
the GPU module. This is useful as `gpu::ObjectAttr` contains a property dict
for metadata, and passing the module allows extracting things like the symbol
table and adding it to the property dict.
---
 mlir/include/mlir/Dialect/GPU/IR/CompilationAttrInterfaces.td | 3 ++-
 mlir/lib/Target/LLVM/NVVM/Target.cpp                          | 4 ++--
 mlir/lib/Target/LLVM/ROCDL/Target.cpp                         | 4 ++--
 mlir/lib/Target/SPIRV/Target.cpp                              | 4 ++--
 4 files changed, 8 insertions(+), 7 deletions(-)

diff --git a/mlir/include/mlir/Dialect/GPU/IR/CompilationAttrInterfaces.td b/mlir/include/mlir/Dialect/GPU/IR/CompilationAttrInterfaces.td
index 1607304803942..3ee4b215f7d64 100644
--- a/mlir/include/mlir/Dialect/GPU/IR/CompilationAttrInterfaces.td
+++ b/mlir/include/mlir/Dialect/GPU/IR/CompilationAttrInterfaces.td
@@ -46,7 +46,8 @@ def GPUTargetAttrInterface : AttrInterface<"TargetAttrInterface"> {
         meant to be used for passing additional options that are not in the
         attribute.
       }], "Attribute", "createObject",
-        (ins "const SmallVector<char, 0>&":$object,
+        (ins "::mlir::Operation*":$module,
+             "const SmallVector<char, 0>&":$object,
              "const gpu::TargetOptions&":$options)>
   ];
 }
diff --git a/mlir/lib/Target/LLVM/NVVM/Target.cpp b/mlir/lib/Target/LLVM/NVVM/Target.cpp
index acb903aa37caf..21e246a020da4 100644
--- a/mlir/lib/Target/LLVM/NVVM/Target.cpp
+++ b/mlir/lib/Target/LLVM/NVVM/Target.cpp
@@ -49,7 +49,7 @@ class NVVMTargetAttrImpl
   serializeToObject(Attribute attribute, Operation *module,
                     const gpu::TargetOptions &options) const;
 
-  Attribute createObject(Attribute attribute,
+  Attribute createObject(Attribute attribute, Operation *module,
                          const SmallVector<char, 0> &object,
                          const gpu::TargetOptions &options) const;
 };
@@ -593,7 +593,7 @@ NVVMTargetAttrImpl::serializeToObject(Attribute attribute, Operation *module,
 }
 
 Attribute
-NVVMTargetAttrImpl::createObject(Attribute attribute,
+NVVMTargetAttrImpl::createObject(Attribute attribute, Operation *module,
                                  const SmallVector<char, 0> &object,
                                  const gpu::TargetOptions &options) const {
   auto target = cast<NVVMTargetAttr>(attribute);
diff --git a/mlir/lib/Target/LLVM/ROCDL/Target.cpp b/mlir/lib/Target/LLVM/ROCDL/Target.cpp
index cc13e5b7436ea..cfa8a7c621d71 100644
--- a/mlir/lib/Target/LLVM/ROCDL/Target.cpp
+++ b/mlir/lib/Target/LLVM/ROCDL/Target.cpp
@@ -62,7 +62,7 @@ class ROCDLTargetAttrImpl
   serializeToObject(Attribute attribute, Operation *module,
                     const gpu::TargetOptions &options) const;
 
-  Attribute createObject(Attribute attribute,
+  Attribute createObject(Attribute attribute, Operation *module,
                          const SmallVector<char, 0> &object,
                          const gpu::TargetOptions &options) const;
 };
@@ -473,7 +473,7 @@ std::optional<SmallVector<char, 0>> ROCDLTargetAttrImpl::serializeToObject(
 }
 
 Attribute
-ROCDLTargetAttrImpl::createObject(Attribute attribute,
+ROCDLTargetAttrImpl::createObject(Attribute attribute, Operation *module,
                                   const SmallVector<char, 0> &object,
                                   const gpu::TargetOptions &options) const {
   gpu::CompilationTarget format = options.getCompilationTarget();
diff --git a/mlir/lib/Target/SPIRV/Target.cpp b/mlir/lib/Target/SPIRV/Target.cpp
index 4c416abe71cac..d48548bf9709c 100644
--- a/mlir/lib/Target/SPIRV/Target.cpp
+++ b/mlir/lib/Target/SPIRV/Target.cpp
@@ -34,7 +34,7 @@ class SPIRVTargetAttrImpl
   serializeToObject(Attribute attribute, Operation *module,
                     const gpu::TargetOptions &options) const;
 
-  Attribute createObject(Attribute attribute,
+  Attribute createObject(Attribute attribute, Operation *module,
                          const SmallVector<char, 0> &object,
                          const gpu::TargetOptions &options) const;
 };
@@ -89,7 +89,7 @@ std::optional<SmallVector<char, 0>> SPIRVTargetAttrImpl::serializeToObject(
 
 // Prepare Attribute for gpu.binary with serialized kernel object
 Attribute
-SPIRVTargetAttrImpl::createObject(Attribute attribute,
+SPIRVTargetAttrImpl::createObject(Attribute attribute, Operation *module,
                                   const SmallVector<char, 0> &object,
                                   const gpu::TargetOptions &options) const {
   gpu::CompilationTarget format = options.getCompilationTarget();

>From 596875cb9f3206fc918176f763ad60f722da9b5d Mon Sep 17 00:00:00 2001
From: Fabian Mora <fmora.dev at gmail.com>
Date: Sun, 9 Jun 2024 20:22:53 +0000
Subject: [PATCH 2/2] Fix missing arg

---
 mlir/lib/Dialect/GPU/Transforms/ModuleToBinary.cpp | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/mlir/lib/Dialect/GPU/Transforms/ModuleToBinary.cpp b/mlir/lib/Dialect/GPU/Transforms/ModuleToBinary.cpp
index 1e7596e8cc4af..ba9df3d00b324 100644
--- a/mlir/lib/Dialect/GPU/Transforms/ModuleToBinary.cpp
+++ b/mlir/lib/Dialect/GPU/Transforms/ModuleToBinary.cpp
@@ -114,7 +114,8 @@ LogicalResult moduleSerializer(GPUModuleOp op,
       return failure();
     }
 
-    Attribute object = target.createObject(*serializedModule, targetOptions);
+    Attribute object =
+        target.createObject(op, *serializedModule, targetOptions);
     if (!object) {
       op.emitError("An error happened while creating the object.");
       return failure();