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

llvmlistbot at llvm.org llvmlistbot at llvm.org
Tue Aug 27 08:05:09 PDT 2024


Author: Fabian Mora
Date: 2024-08-27T11:05:04-04:00
New Revision: fd36a7b9443921a4aa571549543244c5f1c9b6f8

URL: https://github.com/llvm/llvm-project/commit/fd36a7b9443921a4aa571549543244c5f1c9b6f8
DIFF: https://github.com/llvm/llvm-project/commit/fd36a7b9443921a4aa571549543244c5f1c9b6f8.diff

LOG: [mlir][gpu] Pass GPU module to `TargetAttrInterface::createObject`. (#94910)

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, hence the module can be used for extracting
things like the symbol table and adding it to the property dict.

---------

Co-authored-by: Oleksandr "Alex" Zinenko <ftynse at gmail.com>

Added: 
    

Modified: 
    mlir/include/mlir/Dialect/GPU/IR/CompilationAttrInterfaces.td
    mlir/lib/Dialect/GPU/Transforms/ModuleToBinary.cpp
    mlir/lib/Target/LLVM/NVVM/Target.cpp
    mlir/lib/Target/LLVM/ROCDL/Target.cpp
    mlir/lib/Target/SPIRV/Target.cpp
    mlir/unittests/Target/LLVM/SerializeToLLVMBitcode.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Dialect/GPU/IR/CompilationAttrInterfaces.td b/mlir/include/mlir/Dialect/GPU/IR/CompilationAttrInterfaces.td
index 43ca5337067caa..3d73d00ecfdd7e 100644
--- a/mlir/include/mlir/Dialect/GPU/IR/CompilationAttrInterfaces.td
+++ b/mlir/include/mlir/Dialect/GPU/IR/CompilationAttrInterfaces.td
@@ -47,8 +47,9 @@ def GPUTargetAttrInterface : AttrInterface<"TargetAttrInterface"> {
         meant to be used for passing additional options that are not in the
         attribute.
       }], "::mlir::Attribute", "createObject",
-        (ins "const ::llvm::SmallVector<char, 0>&":$object,
-             "const ::mlir::gpu::TargetOptions&":$options)>
+        (ins "::mlir::Operation *":$module,
+             "const ::llvm::SmallVector<char, 0> &":$object,
+             "const ::mlir::gpu::TargetOptions &":$options)>
   ];
 }
 

diff  --git a/mlir/lib/Dialect/GPU/Transforms/ModuleToBinary.cpp b/mlir/lib/Dialect/GPU/Transforms/ModuleToBinary.cpp
index adae3bef763ff7..86a3b4780e88ce 100644
--- a/mlir/lib/Dialect/GPU/Transforms/ModuleToBinary.cpp
+++ b/mlir/lib/Dialect/GPU/Transforms/ModuleToBinary.cpp
@@ -99,7 +99,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();

diff  --git a/mlir/lib/Target/LLVM/NVVM/Target.cpp b/mlir/lib/Target/LLVM/NVVM/Target.cpp
index e608d26e8d2ece..a75b7f92ed8dc3 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;
 };
@@ -591,7 +591,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 4d23f987eb05ef..e32a0c7e14e85c 100644
--- a/mlir/lib/Target/LLVM/ROCDL/Target.cpp
+++ b/mlir/lib/Target/LLVM/ROCDL/Target.cpp
@@ -59,7 +59,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;
 };
@@ -500,7 +500,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 4c416abe71cac9..d48548bf9709c0 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();

diff  --git a/mlir/unittests/Target/LLVM/SerializeToLLVMBitcode.cpp b/mlir/unittests/Target/LLVM/SerializeToLLVMBitcode.cpp
index 96828828fe3f46..37dbfe62036871 100644
--- a/mlir/unittests/Target/LLVM/SerializeToLLVMBitcode.cpp
+++ b/mlir/unittests/Target/LLVM/SerializeToLLVMBitcode.cpp
@@ -6,6 +6,8 @@
 //
 //===----------------------------------------------------------------------===//
 
+#include "mlir/Dialect/GPU/IR/GPUDialect.h"
+#include "mlir/IR/BuiltinDialect.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/MLIRContext.h"
 #include "mlir/Parser/Parser.h"
@@ -30,24 +32,46 @@ using namespace mlir;
 #define SKIP_WITHOUT_NATIVE(x) x
 #endif
 
+namespace {
+// Dummy interface for testing.
+class TargetAttrImpl
+    : public gpu::TargetAttrInterface::FallbackModel<TargetAttrImpl> {
+public:
+  std::optional<SmallVector<char, 0>>
+  serializeToObject(Attribute attribute, Operation *module,
+                    const gpu::TargetOptions &options) const;
+
+  Attribute createObject(Attribute attribute, Operation *module,
+                         const SmallVector<char, 0> &object,
+                         const gpu::TargetOptions &options) const;
+};
+} // namespace
+
 class MLIRTargetLLVM : public ::testing::Test {
 protected:
   void SetUp() override {
     llvm::InitializeNativeTarget();
     llvm::InitializeNativeTargetAsmPrinter();
+    registry.addExtension(+[](MLIRContext *ctx, BuiltinDialect *dialect) {
+      IntegerAttr::attachInterface<TargetAttrImpl>(*ctx);
+    });
+    registerBuiltinDialectTranslation(registry);
+    registerLLVMDialectTranslation(registry);
+    registry.insert<gpu::GPUDialect>();
   }
-};
 
-TEST_F(MLIRTargetLLVM, SKIP_WITHOUT_NATIVE(SerializeToLLVMBitcode)) {
+  // Dialect registry.
+  DialectRegistry registry;
+
+  // MLIR module used for the tests.
   std::string moduleStr = R"mlir(
   llvm.func @foo(%arg0 : i32) {
     llvm.return
   }
   )mlir";
+};
 
-  DialectRegistry registry;
-  registerBuiltinDialectTranslation(registry);
-  registerLLVMDialectTranslation(registry);
+TEST_F(MLIRTargetLLVM, SKIP_WITHOUT_NATIVE(SerializeToLLVMBitcode)) {
   MLIRContext context(registry);
 
   OwningOpRef<ModuleOp> module =
@@ -74,3 +98,52 @@ TEST_F(MLIRTargetLLVM, SKIP_WITHOUT_NATIVE(SerializeToLLVMBitcode)) {
   // Check that it has a function named `foo`.
   ASSERT_TRUE((*llvmModule)->getFunction("foo") != nullptr);
 }
+
+std::optional<SmallVector<char, 0>>
+TargetAttrImpl::serializeToObject(Attribute attribute, Operation *module,
+                                  const gpu::TargetOptions &options) const {
+  module->setAttr("serialize_attr", UnitAttr::get(module->getContext()));
+  std::string targetTriple = llvm::sys::getProcessTriple();
+  LLVM::ModuleToObject serializer(*module, targetTriple, "", "");
+  return serializer.run();
+}
+
+Attribute
+TargetAttrImpl::createObject(Attribute attribute, Operation *module,
+                             const SmallVector<char, 0> &object,
+                             const gpu::TargetOptions &options) const {
+  return gpu::ObjectAttr::get(
+      module->getContext(), attribute, gpu::CompilationTarget::Offload,
+      StringAttr::get(module->getContext(),
+                      StringRef(object.data(), object.size())),
+      module->getAttrDictionary());
+}
+
+TEST_F(MLIRTargetLLVM, SKIP_WITHOUT_NATIVE(TargetAttrAPI)) {
+  MLIRContext context(registry);
+  context.loadAllAvailableDialects();
+
+  OwningOpRef<ModuleOp> module =
+      parseSourceString<ModuleOp>(moduleStr, &context);
+  ASSERT_TRUE(!!module);
+  Builder builder(&context);
+  IntegerAttr target = builder.getI32IntegerAttr(0);
+  auto targetAttr = dyn_cast<gpu::TargetAttrInterface>(target);
+  // Check the attribute holds the interface.
+  ASSERT_TRUE(!!targetAttr);
+  gpu::TargetOptions opts;
+  std::optional<SmallVector<char, 0>> serializedBinary =
+      targetAttr.serializeToObject(*module, opts);
+  // Check the serialized string.
+  ASSERT_TRUE(!!serializedBinary);
+  ASSERT_TRUE(!serializedBinary->empty());
+  // Create the object attribute.
+  auto object = cast<gpu::ObjectAttr>(
+      targetAttr.createObject(*module, *serializedBinary, opts));
+  // Check the object has properties.
+  DictionaryAttr properties = object.getProperties();
+  ASSERT_TRUE(!!properties);
+  // Check that it contains the attribute added to the module in
+  // `serializeToObject`.
+  ASSERT_TRUE(properties.contains("serialize_attr"));
+}


        


More information about the Mlir-commits mailing list