[mlir][LLVM] Add !invariant.group metadata to llvm.load and llvm.store (#115723)

This patch adds support for the `!invariant.group` metadata to the
`llvm.load` and the `llvm.store` operation.
This commit is contained in:
Sirui Mu 2024-11-13 10:54:34 +08:00 committed by GitHub
parent 37143fe27e
commit e887f8290d
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
6 changed files with 77 additions and 5 deletions

View File

@ -229,6 +229,13 @@ class LLVM_MemOpPatterns {
inst->setMetadata(llvm::LLVMContext::MD_nontemporal, metadata);
}
}];
code setInvariantGroupCode = [{
if ($invariantGroup) {
llvm::MDNode *metadata = llvm::MDNode::get(inst->getContext(),
std::nullopt);
inst->setMetadata(llvm::LLVMContext::MD_invariant_group, metadata);
}
}];
code setAccessGroupsMetadataCode = [{
moduleTranslation.setAccessGroupsMetadata(op, inst);
}];

View File

@ -350,6 +350,7 @@ def LLVM_LoadOp : LLVM_MemAccessOpBase<"load",
UnitAttr:$volatile_,
UnitAttr:$nontemporal,
UnitAttr:$invariant,
UnitAttr:$invariantGroup,
DefaultValuedAttr<
AtomicOrdering, "AtomicOrdering::not_atomic">:$ordering,
OptionalAttr<StrAttr>:$syncscope);
@ -385,6 +386,7 @@ def LLVM_LoadOp : LLVM_MemAccessOpBase<"load",
(`volatile` $volatile_^)? $addr
(`atomic` (`syncscope` `(` $syncscope^ `)`)? $ordering^)?
(`invariant` $invariant^)?
(`invariant_group` $invariantGroup^)?
attr-dict `:` qualified(type($addr)) `->` type($res)
}];
string llvmBuilder = [{
@ -398,6 +400,7 @@ def LLVM_LoadOp : LLVM_MemAccessOpBase<"load",
# setSyncScopeCode
# setAlignmentCode
# setNonTemporalMetadataCode
# setInvariantGroupCode
# setAccessGroupsMetadataCode
# setAliasAnalysisMetadataCode;
string mlirBuilder = [{
@ -407,6 +410,7 @@ def LLVM_LoadOp : LLVM_MemAccessOpBase<"load",
alignment, loadInst->isVolatile(),
loadInst->hasMetadata(llvm::LLVMContext::MD_nontemporal),
loadInst->hasMetadata(llvm::LLVMContext::MD_invariant_load),
loadInst->hasMetadata(llvm::LLVMContext::MD_invariant_group),
convertAtomicOrderingFromLLVM(loadInst->getOrdering()),
getLLVMSyncScope(loadInst));
}];
@ -414,6 +418,7 @@ def LLVM_LoadOp : LLVM_MemAccessOpBase<"load",
OpBuilder<(ins "Type":$type, "Value":$addr,
CArg<"unsigned", "0">:$alignment, CArg<"bool", "false">:$isVolatile,
CArg<"bool", "false">:$isNonTemporal, CArg<"bool", "false">:$isInvariant,
CArg<"bool", "false">:$isInvariantGroup,
CArg<"AtomicOrdering", "AtomicOrdering::not_atomic">:$ordering,
CArg<"StringRef", "StringRef()">:$syncscope)>
];
@ -430,6 +435,7 @@ def LLVM_StoreOp : LLVM_MemAccessOpBase<"store",
OptionalAttr<I64Attr>:$alignment,
UnitAttr:$volatile_,
UnitAttr:$nontemporal,
UnitAttr:$invariantGroup,
DefaultValuedAttr<
AtomicOrdering, "AtomicOrdering::not_atomic">:$ordering,
OptionalAttr<StrAttr>:$syncscope);
@ -463,6 +469,7 @@ def LLVM_StoreOp : LLVM_MemAccessOpBase<"store",
let assemblyFormat = [{
(`volatile` $volatile_^)? $value `,` $addr
(`atomic` (`syncscope` `(` $syncscope^ `)`)? $ordering^)?
(`invariant_group` $invariantGroup^)?
attr-dict `:` type($value) `,` qualified(type($addr))
}];
string llvmBuilder = [{
@ -471,6 +478,7 @@ def LLVM_StoreOp : LLVM_MemAccessOpBase<"store",
# setSyncScopeCode
# setAlignmentCode
# setNonTemporalMetadataCode
# setInvariantGroupCode
# setAccessGroupsMetadataCode
# setAliasAnalysisMetadataCode;
string mlirBuilder = [{
@ -479,6 +487,7 @@ def LLVM_StoreOp : LLVM_MemAccessOpBase<"store",
$_op = $_builder.create<LLVM::StoreOp>($_location, $value, $addr,
alignment, storeInst->isVolatile(),
storeInst->hasMetadata(llvm::LLVMContext::MD_nontemporal),
storeInst->hasMetadata(llvm::LLVMContext::MD_invariant_group),
convertAtomicOrderingFromLLVM(storeInst->getOrdering()),
getLLVMSyncScope(storeInst));
}];
@ -486,6 +495,7 @@ def LLVM_StoreOp : LLVM_MemAccessOpBase<"store",
OpBuilder<(ins "Value":$value, "Value":$addr,
CArg<"unsigned", "0">:$alignment, CArg<"bool", "false">:$isVolatile,
CArg<"bool", "false">:$isNonTemporal,
CArg<"bool", "false">:$isInvariantGroup,
CArg<"AtomicOrdering", "AtomicOrdering::not_atomic">:$ordering,
CArg<"StringRef", "StringRef()">:$syncscope)>
];

View File

@ -944,11 +944,11 @@ LogicalResult LoadOp::verify() {
void LoadOp::build(OpBuilder &builder, OperationState &state, Type type,
Value addr, unsigned alignment, bool isVolatile,
bool isNonTemporal, bool isInvariant,
bool isNonTemporal, bool isInvariant, bool isInvariantGroup,
AtomicOrdering ordering, StringRef syncscope) {
build(builder, state, type, addr,
alignment ? builder.getI64IntegerAttr(alignment) : nullptr, isVolatile,
isNonTemporal, isInvariant, ordering,
isNonTemporal, isInvariant, isInvariantGroup, ordering,
syncscope.empty() ? nullptr : builder.getStringAttr(syncscope),
/*access_groups=*/nullptr,
/*alias_scopes=*/nullptr, /*noalias_scopes=*/nullptr,
@ -983,11 +983,11 @@ LogicalResult StoreOp::verify() {
void StoreOp::build(OpBuilder &builder, OperationState &state, Value value,
Value addr, unsigned alignment, bool isVolatile,
bool isNonTemporal, AtomicOrdering ordering,
StringRef syncscope) {
bool isNonTemporal, bool isInvariantGroup,
AtomicOrdering ordering, StringRef syncscope) {
build(builder, state, value, addr,
alignment ? builder.getI64IntegerAttr(alignment) : nullptr, isVolatile,
isNonTemporal, ordering,
isNonTemporal, isInvariantGroup, ordering,
syncscope.empty() ? nullptr : builder.getStringAttr(syncscope),
/*access_groups=*/nullptr,
/*alias_scopes=*/nullptr, /*noalias_scopes=*/nullptr, /*tbaa=*/nullptr);

View File

@ -486,6 +486,20 @@ func.func @invariant_load(%ptr : !llvm.ptr) -> i32 {
func.return %0 : i32
}
// CHECK-LABEL: @invariant_group_load
func.func @invariant_group_load(%ptr : !llvm.ptr) -> i32 {
// CHECK: llvm.load %{{.+}} invariant_group {alignment = 4 : i64} : !llvm.ptr -> i32
%0 = llvm.load %ptr invariant_group {alignment = 4 : i64} : !llvm.ptr -> i32
func.return %0 : i32
}
// CHECK-LABEL: @invariant_group_store
func.func @invariant_group_store(%val: i32, %ptr : !llvm.ptr) {
// CHECK: llvm.store %{{.+}}, %{{.+}} invariant_group : i32, !llvm.ptr
llvm.store %val, %ptr invariant_group : i32, !llvm.ptr
func.return
}
llvm.mlir.global external constant @_ZTIi() : !llvm.ptr
llvm.func @bar(!llvm.ptr, !llvm.ptr, !llvm.ptr)
llvm.func @__gxx_personality_v0(...) -> i32

View File

@ -383,6 +383,33 @@ define float @invariant_load(ptr %ptr) {
; // -----
; CHECK-LABEL: @invariant_group_load
; CHECK-SAME: %[[PTR:[a-zA-Z0-9]+]]
define float @invariant_group_load(ptr %ptr) {
; CHECK: %[[VAL:.+]] = llvm.load %[[PTR]] invariant_group {alignment = 4 : i64} : !llvm.ptr -> f32
%1 = load float, ptr %ptr, align 4, !invariant.group !0
; CHECK: llvm.return %[[VAL]]
ret float %1
}
!0 = !{}
; // -----
; CHECK-LABEL: @invariant_group_store
; CHECK-SAME: %[[VAL:[a-zA-Z0-9]+]]
; CHECK-SAME: %[[PTR:[a-zA-Z0-9]+]]
define void @invariant_group_store(float %val, ptr %ptr) {
; CHECK: llvm.store %[[VAL]], %[[PTR]] invariant_group {alignment = 4 : i64} : f32, !llvm.ptr
store float %val, ptr %ptr, align 4, !invariant.group !0
; CHECK: llvm.return
ret void
}
!0 = !{}
; // -----
; CHECK-LABEL: @atomic_load_store
; CHECK-SAME: %[[PTR:[a-zA-Z0-9]+]]
define void @atomic_load_store(ptr %ptr) {

View File

@ -1967,6 +1967,20 @@ llvm.func @nontemporal_store_and_load(%ptr : !llvm.ptr) -> i32 {
// -----
// Check that invariant group attribute is exported as metadata node.
llvm.func @nontemporal_store_and_load(%ptr : !llvm.ptr) -> i32 {
%val = llvm.mlir.constant(42 : i32) : i32
// CHECK: store i32 42, ptr %{{.*}} !invariant.group ![[NODE:[0-9]+]]
llvm.store %val, %ptr invariant_group : i32, !llvm.ptr
// CHECK: %{{.*}} = load i32, ptr %{{.*}} !invariant.group ![[NODE]]
%1 = llvm.load %ptr invariant_group : !llvm.ptr -> i32
llvm.return %1 : i32
}
// CHECK: ![[NODE]] = !{}
// -----
llvm.func @atomic_store_and_load(%ptr : !llvm.ptr) {
// CHECK: load atomic
// CHECK-SAME: acquire, align 4