Skip to content

Commit fd36a7b

Browse files
fabianmcgftynse
fabianmcg
and
ftynse
authored
[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 <[email protected]>
1 parent d4ffccf commit fd36a7b

File tree

6 files changed

+89
-14
lines changed

6 files changed

+89
-14
lines changed

‎mlir/include/mlir/Dialect/GPU/IR/CompilationAttrInterfaces.td

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -47,8 +47,9 @@ def GPUTargetAttrInterface : AttrInterface<"TargetAttrInterface"> {
4747
meant to be used for passing additional options that are not in the
4848
attribute.
4949
}], "::mlir::Attribute", "createObject",
50-
(ins "const ::llvm::SmallVector<char, 0>&":$object,
51-
"const ::mlir::gpu::TargetOptions&":$options)>
50+
(ins "::mlir::Operation *":$module,
51+
"const ::llvm::SmallVector<char, 0> &":$object,
52+
"const ::mlir::gpu::TargetOptions &":$options)>
5253
];
5354
}
5455

‎mlir/lib/Dialect/GPU/Transforms/ModuleToBinary.cpp

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -99,7 +99,8 @@ LogicalResult moduleSerializer(GPUModuleOp op,
9999
return failure();
100100
}
101101

102-
Attribute object = target.createObject(*serializedModule, targetOptions);
102+
Attribute object =
103+
target.createObject(op, *serializedModule, targetOptions);
103104
if (!object) {
104105
op.emitError("An error happened while creating the object.");
105106
return failure();

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -49,7 +49,7 @@ class NVVMTargetAttrImpl
4949
serializeToObject(Attribute attribute, Operation *module,
5050
const gpu::TargetOptions &options) const;
5151

52-
Attribute createObject(Attribute attribute,
52+
Attribute createObject(Attribute attribute, Operation *module,
5353
const SmallVector<char, 0> &object,
5454
const gpu::TargetOptions &options) const;
5555
};
@@ -591,7 +591,7 @@ NVVMTargetAttrImpl::serializeToObject(Attribute attribute, Operation *module,
591591
}
592592

593593
Attribute
594-
NVVMTargetAttrImpl::createObject(Attribute attribute,
594+
NVVMTargetAttrImpl::createObject(Attribute attribute, Operation *module,
595595
const SmallVector<char, 0> &object,
596596
const gpu::TargetOptions &options) const {
597597
auto target = cast<NVVMTargetAttr>(attribute);

‎mlir/lib/Target/LLVM/ROCDL/Target.cpp

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -59,7 +59,7 @@ class ROCDLTargetAttrImpl
5959
serializeToObject(Attribute attribute, Operation *module,
6060
const gpu::TargetOptions &options) const;
6161

62-
Attribute createObject(Attribute attribute,
62+
Attribute createObject(Attribute attribute, Operation *module,
6363
const SmallVector<char, 0> &object,
6464
const gpu::TargetOptions &options) const;
6565
};
@@ -500,7 +500,7 @@ std::optional<SmallVector<char, 0>> ROCDLTargetAttrImpl::serializeToObject(
500500
}
501501

502502
Attribute
503-
ROCDLTargetAttrImpl::createObject(Attribute attribute,
503+
ROCDLTargetAttrImpl::createObject(Attribute attribute, Operation *module,
504504
const SmallVector<char, 0> &object,
505505
const gpu::TargetOptions &options) const {
506506
gpu::CompilationTarget format = options.getCompilationTarget();

‎mlir/lib/Target/SPIRV/Target.cpp

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -34,7 +34,7 @@ class SPIRVTargetAttrImpl
3434
serializeToObject(Attribute attribute, Operation *module,
3535
const gpu::TargetOptions &options) const;
3636

37-
Attribute createObject(Attribute attribute,
37+
Attribute createObject(Attribute attribute, Operation *module,
3838
const SmallVector<char, 0> &object,
3939
const gpu::TargetOptions &options) const;
4040
};
@@ -89,7 +89,7 @@ std::optional<SmallVector<char, 0>> SPIRVTargetAttrImpl::serializeToObject(
8989

9090
// Prepare Attribute for gpu.binary with serialized kernel object
9191
Attribute
92-
SPIRVTargetAttrImpl::createObject(Attribute attribute,
92+
SPIRVTargetAttrImpl::createObject(Attribute attribute, Operation *module,
9393
const SmallVector<char, 0> &object,
9494
const gpu::TargetOptions &options) const {
9595
gpu::CompilationTarget format = options.getCompilationTarget();

‎mlir/unittests/Target/LLVM/SerializeToLLVMBitcode.cpp

Lines changed: 78 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,8 @@
66
//
77
//===----------------------------------------------------------------------===//
88

9+
#include "mlir/Dialect/GPU/IR/GPUDialect.h"
10+
#include "mlir/IR/BuiltinDialect.h"
911
#include "mlir/IR/BuiltinOps.h"
1012
#include "mlir/IR/MLIRContext.h"
1113
#include "mlir/Parser/Parser.h"
@@ -30,24 +32,46 @@ using namespace mlir;
3032
#define SKIP_WITHOUT_NATIVE(x) x
3133
#endif
3234

35+
namespace {
36+
// Dummy interface for testing.
37+
class TargetAttrImpl
38+
: public gpu::TargetAttrInterface::FallbackModel<TargetAttrImpl> {
39+
public:
40+
std::optional<SmallVector<char, 0>>
41+
serializeToObject(Attribute attribute, Operation *module,
42+
const gpu::TargetOptions &options) const;
43+
44+
Attribute createObject(Attribute attribute, Operation *module,
45+
const SmallVector<char, 0> &object,
46+
const gpu::TargetOptions &options) const;
47+
};
48+
} // namespace
49+
3350
class MLIRTargetLLVM : public ::testing::Test {
3451
protected:
3552
void SetUp() override {
3653
llvm::InitializeNativeTarget();
3754
llvm::InitializeNativeTargetAsmPrinter();
55+
registry.addExtension(+[](MLIRContext *ctx, BuiltinDialect *dialect) {
56+
IntegerAttr::attachInterface<TargetAttrImpl>(*ctx);
57+
});
58+
registerBuiltinDialectTranslation(registry);
59+
registerLLVMDialectTranslation(registry);
60+
registry.insert<gpu::GPUDialect>();
3861
}
39-
};
4062

41-
TEST_F(MLIRTargetLLVM, SKIP_WITHOUT_NATIVE(SerializeToLLVMBitcode)) {
63+
// Dialect registry.
64+
DialectRegistry registry;
65+
66+
// MLIR module used for the tests.
4267
std::string moduleStr = R"mlir(
4368
llvm.func @foo(%arg0 : i32) {
4469
llvm.return
4570
}
4671
)mlir";
72+
};
4773

48-
DialectRegistry registry;
49-
registerBuiltinDialectTranslation(registry);
50-
registerLLVMDialectTranslation(registry);
74+
TEST_F(MLIRTargetLLVM, SKIP_WITHOUT_NATIVE(SerializeToLLVMBitcode)) {
5175
MLIRContext context(registry);
5276

5377
OwningOpRef<ModuleOp> module =
@@ -74,3 +98,52 @@ TEST_F(MLIRTargetLLVM, SKIP_WITHOUT_NATIVE(SerializeToLLVMBitcode)) {
7498
// Check that it has a function named `foo`.
7599
ASSERT_TRUE((*llvmModule)->getFunction("foo") != nullptr);
76100
}
101+
102+
std::optional<SmallVector<char, 0>>
103+
TargetAttrImpl::serializeToObject(Attribute attribute, Operation *module,
104+
const gpu::TargetOptions &options) const {
105+
module->setAttr("serialize_attr", UnitAttr::get(module->getContext()));
106+
std::string targetTriple = llvm::sys::getProcessTriple();
107+
LLVM::ModuleToObject serializer(*module, targetTriple, "", "");
108+
return serializer.run();
109+
}
110+
111+
Attribute
112+
TargetAttrImpl::createObject(Attribute attribute, Operation *module,
113+
const SmallVector<char, 0> &object,
114+
const gpu::TargetOptions &options) const {
115+
return gpu::ObjectAttr::get(
116+
module->getContext(), attribute, gpu::CompilationTarget::Offload,
117+
StringAttr::get(module->getContext(),
118+
StringRef(object.data(), object.size())),
119+
module->getAttrDictionary());
120+
}
121+
122+
TEST_F(MLIRTargetLLVM, SKIP_WITHOUT_NATIVE(TargetAttrAPI)) {
123+
MLIRContext context(registry);
124+
context.loadAllAvailableDialects();
125+
126+
OwningOpRef<ModuleOp> module =
127+
parseSourceString<ModuleOp>(moduleStr, &context);
128+
ASSERT_TRUE(!!module);
129+
Builder builder(&context);
130+
IntegerAttr target = builder.getI32IntegerAttr(0);
131+
auto targetAttr = dyn_cast<gpu::TargetAttrInterface>(target);
132+
// Check the attribute holds the interface.
133+
ASSERT_TRUE(!!targetAttr);
134+
gpu::TargetOptions opts;
135+
std::optional<SmallVector<char, 0>> serializedBinary =
136+
targetAttr.serializeToObject(*module, opts);
137+
// Check the serialized string.
138+
ASSERT_TRUE(!!serializedBinary);
139+
ASSERT_TRUE(!serializedBinary->empty());
140+
// Create the object attribute.
141+
auto object = cast<gpu::ObjectAttr>(
142+
targetAttr.createObject(*module, *serializedBinary, opts));
143+
// Check the object has properties.
144+
DictionaryAttr properties = object.getProperties();
145+
ASSERT_TRUE(!!properties);
146+
// Check that it contains the attribute added to the module in
147+
// `serializeToObject`.
148+
ASSERT_TRUE(properties.contains("serialize_attr"));
149+
}

0 commit comments

Comments
 (0)