[Mlir-commits] [mlir] [mlir][gpu] Pass GPU module to `TargetAttrInterface::createObject`. (PR #94910)
Fabian Mora
llvmlistbot at llvm.org
Sun Jun 9 15:48:24 PDT 2024
https://github.com/fabianmcg updated https://github.com/llvm/llvm-project/pull/94910
>From 887f24d89b00aa07cd4662dcc195359eb70b9055 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/3] [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 582ab91311974..9d812e4119943 100644
--- a/mlir/include/mlir/Dialect/GPU/IR/CompilationAttrInterfaces.td
+++ b/mlir/include/mlir/Dialect/GPU/IR/CompilationAttrInterfaces.td
@@ -47,7 +47,8 @@ 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,
+ (ins "::mlir::Operation*":$module,
+ "const ::llvm::SmallVector<char, 0>&":$object,
"const ::mlir::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 5d7a28134a4d013986a91bdea427552e57395b6d 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/3] 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();
>From 8527ee104f188a99e9c747ed4080decb042a3aa8 Mon Sep 17 00:00:00 2001
From: Fabian Mora <fmora.dev at gmail.com>
Date: Sun, 9 Jun 2024 21:31:38 +0000
Subject: [PATCH 3/3] added tests and passthrough discardable attrs
---
mlir/include/mlir/Dialect/GPU/IR/GPUDialect.h | 1 -
mlir/lib/Target/LLVM/NVVM/Target.cpp | 13 ++++++++-----
mlir/lib/Target/LLVM/ROCDL/Target.cpp | 4 +++-
mlir/lib/Target/SPIRV/Target.cpp | 4 ++--
mlir/test/Dialect/GPU/module-to-binary-nvvm.mlir | 10 ++++++++++
mlir/test/Dialect/GPU/module-to-binary-rocdl.mlir | 10 ++++++++++
mlir/test/Dialect/GPU/module-to-binary-spirv.mlir | 9 +++++++++
7 files changed, 42 insertions(+), 9 deletions(-)
diff --git a/mlir/include/mlir/Dialect/GPU/IR/GPUDialect.h b/mlir/include/mlir/Dialect/GPU/IR/GPUDialect.h
index 96e1935bd0a84..c8f618e8a8c88 100644
--- a/mlir/include/mlir/Dialect/GPU/IR/GPUDialect.h
+++ b/mlir/include/mlir/Dialect/GPU/IR/GPUDialect.h
@@ -202,7 +202,6 @@ class SparseSpGEMMOpHandleType
static constexpr StringLiteral name = "gpu.sparse.spgemmop_handle";
};
-
} // namespace gpu
} // namespace mlir
diff --git a/mlir/lib/Target/LLVM/NVVM/Target.cpp b/mlir/lib/Target/LLVM/NVVM/Target.cpp
index 21e246a020da4..217d56bc779e8 100644
--- a/mlir/lib/Target/LLVM/NVVM/Target.cpp
+++ b/mlir/lib/Target/LLVM/NVVM/Target.cpp
@@ -598,13 +598,16 @@ NVVMTargetAttrImpl::createObject(Attribute attribute, Operation *module,
const gpu::TargetOptions &options) const {
auto target = cast<NVVMTargetAttr>(attribute);
gpu::CompilationTarget format = options.getCompilationTarget();
- DictionaryAttr objectProps;
+ DictionaryAttr objectProps = module->getDiscardableAttrDictionary();
Builder builder(attribute.getContext());
- if (format == gpu::CompilationTarget::Assembly)
- objectProps = builder.getDictionaryAttr(
- {builder.getNamedAttr("O", builder.getI32IntegerAttr(target.getO()))});
+ if (format == gpu::CompilationTarget::Assembly) {
+ SmallVector<NamedAttribute, 1> attrs(objectProps.getValue());
+ attrs.push_back(
+ builder.getNamedAttr("O", builder.getI32IntegerAttr(target.getO())));
+ objectProps = builder.getDictionaryAttr(attrs);
+ }
return builder.getAttr<gpu::ObjectAttr>(
attribute, format,
builder.getStringAttr(StringRef(object.data(), object.size())),
- objectProps);
+ objectProps.empty() ? nullptr : objectProps);
}
diff --git a/mlir/lib/Target/LLVM/ROCDL/Target.cpp b/mlir/lib/Target/LLVM/ROCDL/Target.cpp
index cfa8a7c621d71..c8908bc94408c 100644
--- a/mlir/lib/Target/LLVM/ROCDL/Target.cpp
+++ b/mlir/lib/Target/LLVM/ROCDL/Target.cpp
@@ -477,10 +477,12 @@ ROCDLTargetAttrImpl::createObject(Attribute attribute, Operation *module,
const SmallVector<char, 0> &object,
const gpu::TargetOptions &options) const {
gpu::CompilationTarget format = options.getCompilationTarget();
+ DictionaryAttr objectProps = module->getDiscardableAttrDictionary();
Builder builder(attribute.getContext());
return builder.getAttr<gpu::ObjectAttr>(
attribute,
format > gpu::CompilationTarget::Binary ? gpu::CompilationTarget::Binary
: format,
- builder.getStringAttr(StringRef(object.data(), object.size())), nullptr);
+ builder.getStringAttr(StringRef(object.data(), object.size())),
+ objectProps.empty() ? nullptr : objectProps);
}
diff --git a/mlir/lib/Target/SPIRV/Target.cpp b/mlir/lib/Target/SPIRV/Target.cpp
index d48548bf9709c..700312b7f575c 100644
--- a/mlir/lib/Target/SPIRV/Target.cpp
+++ b/mlir/lib/Target/SPIRV/Target.cpp
@@ -93,10 +93,10 @@ SPIRVTargetAttrImpl::createObject(Attribute attribute, Operation *module,
const SmallVector<char, 0> &object,
const gpu::TargetOptions &options) const {
gpu::CompilationTarget format = options.getCompilationTarget();
- DictionaryAttr objectProps;
+ DictionaryAttr objectProps = module->getDiscardableAttrDictionary();
Builder builder(attribute.getContext());
return builder.getAttr<gpu::ObjectAttr>(
attribute, format,
builder.getStringAttr(StringRef(object.data(), object.size())),
- objectProps);
+ objectProps.empty() ? nullptr : objectProps);
}
diff --git a/mlir/test/Dialect/GPU/module-to-binary-nvvm.mlir b/mlir/test/Dialect/GPU/module-to-binary-nvvm.mlir
index c286c8bc9042f..c3b82971b783b 100644
--- a/mlir/test/Dialect/GPU/module-to-binary-nvvm.mlir
+++ b/mlir/test/Dialect/GPU/module-to-binary-nvvm.mlir
@@ -34,4 +34,14 @@ module attributes {gpu.container_module} {
llvm.return
}
}
+
+ // CHECK-LABEL:gpu.binary @kernel_module4
+ // CHECK-ISA:[#gpu.object<#nvvm.target<flags = {fast}>, properties = {O = 2 : i32, gpu.reg_count = 255 : i32}, assembly = "{{.*}}">]
+ gpu.module @kernel_module4 [#nvvm.target<flags = {fast}>] attributes {gpu.reg_count = 255 : i32} {
+ llvm.func @kernel(%arg0: i32, %arg1: !llvm.ptr,
+ %arg2: !llvm.ptr, %arg3: i64, %arg4: i64,
+ %arg5: i64) attributes {gpu.kernel} {
+ llvm.return
+ }
+ }
}
diff --git a/mlir/test/Dialect/GPU/module-to-binary-rocdl.mlir b/mlir/test/Dialect/GPU/module-to-binary-rocdl.mlir
index 939dbdd4382e7..155d9cd3b314b 100644
--- a/mlir/test/Dialect/GPU/module-to-binary-rocdl.mlir
+++ b/mlir/test/Dialect/GPU/module-to-binary-rocdl.mlir
@@ -22,4 +22,14 @@ module attributes {gpu.container_module} {
llvm.return
}
}
+
+ // CHECK-LABEL:gpu.binary @kernel_module3
+ // CHECK-ISA:[#gpu.object<#rocdl.target<flags = {fast}>, properties = {gpu.reg_count = 255 : i32}, assembly = "{{.*}}">]
+ gpu.module @kernel_module3 [#rocdl.target<flags = {fast}>] attributes {gpu.reg_count = 255 : i32} {
+ llvm.func @kernel(%arg0: i32, %arg1: !llvm.ptr,
+ %arg2: !llvm.ptr, %arg3: i64, %arg4: i64,
+ %arg5: i64) attributes {gpu.kernel} {
+ llvm.return
+ }
+ }
}
diff --git a/mlir/test/Dialect/GPU/module-to-binary-spirv.mlir b/mlir/test/Dialect/GPU/module-to-binary-spirv.mlir
index d62f439279801..06ab7b52dc737 100644
--- a/mlir/test/Dialect/GPU/module-to-binary-spirv.mlir
+++ b/mlir/test/Dialect/GPU/module-to-binary-spirv.mlir
@@ -10,4 +10,13 @@ module attributes {gpu.container_module} {
}
}
}
+ // CHECK-LABEL:gpu.binary @kernel_module1
+ // CHECK:[#gpu.object<#spirv.target_env<#spirv.vce<v1.0, [Int64, Int16, Kernel, Addresses], []>, #spirv.resource_limits<>>, properties = {gpu.reg_count = 255 : i32}, "{{.*}}">]
+ gpu.module @kernel_module1 [#spirv.target_env<#spirv.vce<v1.0, [Int64, Int16, Kernel, Addresses], []>, #spirv.resource_limits<>>] attributes {gpu.reg_count = 255 : i32} {
+ spirv.module @__spv__kernel_module Physical64 OpenCL requires #spirv.vce<v1.0, [Int64, Int16, Kernel, Addresses], []> attributes {spirv.target_env = #spirv.target_env<#spirv.vce<v1.0, [Addresses, Float16Buffer, Int64, Int16, Int8, Kernel, Linkage, Vector16, GenericPointer, Groups, Float16, Float64, AtomicFloat32AddEXT, ExpectAssumeKHR], [SPV_EXT_shader_atomic_float_add, SPV_KHR_expect_assume, SPV_AMD_shader_ballot]>, api=OpenCL, #spirv.resource_limits<>>} {
+ spirv.func @test_kernel(%arg0: !spirv.ptr<!spirv.array<200 x i16>, CrossWorkgroup>, %arg1: !spirv.ptr<!spirv.array<200 x i16>, CrossWorkgroup>, %arg2: !spirv.ptr<!spirv.array<200 x i16>, CrossWorkgroup>) "None" attributes {workgroup_attributions = 0 : i64} {
+ spirv.Return
+ }
+ }
+ }
}
More information about the Mlir-commits
mailing list