Skip to content

Commit 09b79c6

Browse files
LancernLancern
committed
[mlir][LLVM] Add !invariant.group metadata to llvm.load and llvm.store
This patch adds support for the `!invariant.group` metadata to the `llvm.load` and the `llvm.store` operation.
1 parent 89aaf2c commit 09b79c6

File tree

6 files changed

+77
-5
lines changed

6 files changed

+77
-5
lines changed

‎mlir/include/mlir/Dialect/LLVMIR/LLVMOpBase.td‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -229,6 +229,13 @@ class LLVM_MemOpPatterns {
229229
inst->setMetadata(llvm::LLVMContext::MD_nontemporal, metadata);
230230
}
231231
}];
232+
code setInvariantGroupCode = [{
233+
if ($invariantGroup) {
234+
llvm::MDNode *metadata = llvm::MDNode::get(inst->getContext(),
235+
std::nullopt);
236+
inst->setMetadata(llvm::LLVMContext::MD_invariant_group, metadata);
237+
}
238+
}];
232239
code setAccessGroupsMetadataCode = [{
233240
moduleTranslation.setAccessGroupsMetadata(op, inst);
234241
}];

‎mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -351,6 +351,7 @@ def LLVM_LoadOp : LLVM_MemAccessOpBase<"load",
351351
UnitAttr:$volatile_,
352352
UnitAttr:$nontemporal,
353353
UnitAttr:$invariant,
354+
UnitAttr:$invariantGroup,
354355
DefaultValuedAttr<
355356
AtomicOrdering, "AtomicOrdering::not_atomic">:$ordering,
356357
OptionalAttr<StrAttr>:$syncscope);
@@ -386,6 +387,7 @@ def LLVM_LoadOp : LLVM_MemAccessOpBase<"load",
386387
(`volatile` $volatile_^)? $addr
387388
(`atomic` (`syncscope` `(` $syncscope^ `)`)? $ordering^)?
388389
(`invariant` $invariant^)?
390+
(`invariant_group` $invariantGroup^)?
389391
attr-dict `:` qualified(type($addr)) `->` type($res)
390392
}];
391393
string llvmBuilder = [{
@@ -399,6 +401,7 @@ def LLVM_LoadOp : LLVM_MemAccessOpBase<"load",
399401
# setSyncScopeCode
400402
# setAlignmentCode
401403
# setNonTemporalMetadataCode
404+
# setInvariantGroupCode
402405
# setAccessGroupsMetadataCode
403406
# setAliasAnalysisMetadataCode;
404407
string mlirBuilder = [{
@@ -408,13 +411,15 @@ def LLVM_LoadOp : LLVM_MemAccessOpBase<"load",
408411
alignment, loadInst->isVolatile(),
409412
loadInst->hasMetadata(llvm::LLVMContext::MD_nontemporal),
410413
loadInst->hasMetadata(llvm::LLVMContext::MD_invariant_load),
414+
loadInst->hasMetadata(llvm::LLVMContext::MD_invariant_group),
411415
convertAtomicOrderingFromLLVM(loadInst->getOrdering()),
412416
getLLVMSyncScope(loadInst));
413417
}];
414418
let builders = [
415419
OpBuilder<(ins "Type":$type, "Value":$addr,
416420
CArg<"unsigned", "0">:$alignment, CArg<"bool", "false">:$isVolatile,
417421
CArg<"bool", "false">:$isNonTemporal, CArg<"bool", "false">:$isInvariant,
422+
CArg<"bool", "false">:$isInvariantGroup,
418423
CArg<"AtomicOrdering", "AtomicOrdering::not_atomic">:$ordering,
419424
CArg<"StringRef", "StringRef()">:$syncscope)>
420425
];
@@ -431,6 +436,7 @@ def LLVM_StoreOp : LLVM_MemAccessOpBase<"store",
431436
OptionalAttr<I64Attr>:$alignment,
432437
UnitAttr:$volatile_,
433438
UnitAttr:$nontemporal,
439+
UnitAttr:$invariantGroup,
434440
DefaultValuedAttr<
435441
AtomicOrdering, "AtomicOrdering::not_atomic">:$ordering,
436442
OptionalAttr<StrAttr>:$syncscope);
@@ -464,6 +470,7 @@ def LLVM_StoreOp : LLVM_MemAccessOpBase<"store",
464470
let assemblyFormat = [{
465471
(`volatile` $volatile_^)? $value `,` $addr
466472
(`atomic` (`syncscope` `(` $syncscope^ `)`)? $ordering^)?
473+
(`invariant_group` $invariantGroup^)?
467474
attr-dict `:` type($value) `,` qualified(type($addr))
468475
}];
469476
string llvmBuilder = [{
@@ -472,6 +479,7 @@ def LLVM_StoreOp : LLVM_MemAccessOpBase<"store",
472479
# setSyncScopeCode
473480
# setAlignmentCode
474481
# setNonTemporalMetadataCode
482+
# setInvariantGroupCode
475483
# setAccessGroupsMetadataCode
476484
# setAliasAnalysisMetadataCode;
477485
string mlirBuilder = [{
@@ -480,13 +488,15 @@ def LLVM_StoreOp : LLVM_MemAccessOpBase<"store",
480488
$_op = $_builder.create<LLVM::StoreOp>($_location, $value, $addr,
481489
alignment, storeInst->isVolatile(),
482490
storeInst->hasMetadata(llvm::LLVMContext::MD_nontemporal),
491+
storeInst->hasMetadata(llvm::LLVMContext::MD_invariant_group),
483492
convertAtomicOrderingFromLLVM(storeInst->getOrdering()),
484493
getLLVMSyncScope(storeInst));
485494
}];
486495
let builders = [
487496
OpBuilder<(ins "Value":$value, "Value":$addr,
488497
CArg<"unsigned", "0">:$alignment, CArg<"bool", "false">:$isVolatile,
489498
CArg<"bool", "false">:$isNonTemporal,
499+
CArg<"bool", "false">:$isInvariantGroup,
490500
CArg<"AtomicOrdering", "AtomicOrdering::not_atomic">:$ordering,
491501
CArg<"StringRef", "StringRef()">:$syncscope)>
492502
];

‎mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -952,11 +952,11 @@ LogicalResult LoadOp::verify() {
952952

953953
void LoadOp::build(OpBuilder &builder, OperationState &state, Type type,
954954
Value addr, unsigned alignment, bool isVolatile,
955-
bool isNonTemporal, bool isInvariant,
955+
bool isNonTemporal, bool isInvariant, bool isInvariantGroup,
956956
AtomicOrdering ordering, StringRef syncscope) {
957957
build(builder, state, type, addr,
958958
alignment ? builder.getI64IntegerAttr(alignment) : nullptr, isVolatile,
959-
isNonTemporal, isInvariant, ordering,
959+
isNonTemporal, isInvariant, isInvariantGroup, ordering,
960960
syncscope.empty() ? nullptr : builder.getStringAttr(syncscope),
961961
/*access_groups=*/nullptr,
962962
/*alias_scopes=*/nullptr, /*noalias_scopes=*/nullptr,
@@ -991,11 +991,11 @@ LogicalResult StoreOp::verify() {
991991

992992
void StoreOp::build(OpBuilder &builder, OperationState &state, Value value,
993993
Value addr, unsigned alignment, bool isVolatile,
994-
bool isNonTemporal, AtomicOrdering ordering,
995-
StringRef syncscope) {
994+
bool isNonTemporal, bool isInvariantGroup,
995+
AtomicOrdering ordering, StringRef syncscope) {
996996
build(builder, state, value, addr,
997997
alignment ? builder.getI64IntegerAttr(alignment) : nullptr, isVolatile,
998-
isNonTemporal, ordering,
998+
isNonTemporal, isInvariantGroup, ordering,
999999
syncscope.empty() ? nullptr : builder.getStringAttr(syncscope),
10001000
/*access_groups=*/nullptr,
10011001
/*alias_scopes=*/nullptr, /*noalias_scopes=*/nullptr, /*tbaa=*/nullptr);

‎mlir/test/Dialect/LLVMIR/roundtrip.mlir‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -469,6 +469,20 @@ func.func @invariant_load(%ptr : !llvm.ptr) -> i32 {
469469
func.return %0 : i32
470470
}
471471

472+
// CHECK-LABEL: @invariant_group_load
473+
func.func @invariant_group_load(%ptr : !llvm.ptr) -> i32 {
474+
// CHECK: llvm.load %{{.+}} invariant_group {alignment = 4 : i64} : !llvm.ptr -> i32
475+
%0 = llvm.load %ptr invariant_group {alignment = 4 : i64} : !llvm.ptr -> i32
476+
func.return %0 : i32
477+
}
478+
479+
// CHECK-LABEL: @invariant_group_store
480+
func.func @invariant_group_store(%val: i32, %ptr : !llvm.ptr) {
481+
// CHECK: llvm.store %{{.+}}, %{{.+}} invariant_group : i32, !llvm.ptr
482+
llvm.store %val, %ptr invariant_group : i32, !llvm.ptr
483+
func.return
484+
}
485+
472486
llvm.mlir.global external constant @_ZTIi() : !llvm.ptr
473487
llvm.func @bar(!llvm.ptr, !llvm.ptr, !llvm.ptr)
474488
llvm.func @__gxx_personality_v0(...) -> i32

‎mlir/test/Target/LLVMIR/Import/instructions.ll‎

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -383,6 +383,33 @@ define float @invariant_load(ptr %ptr) {
383383

384384
; // -----
385385

386+
; CHECK-LABEL: @invariant_group_load
387+
; CHECK-SAME: %[[PTR:[a-zA-Z0-9]+]]
388+
define float @invariant_group_load(ptr %ptr) {
389+
; CHECK: %[[VAL:.+]] = llvm.load %[[PTR]] invariant_group {alignment = 4 : i64} : !llvm.ptr -> f32
390+
%1 = load float, ptr %ptr, align 4, !invariant.group !0
391+
; CHECK: llvm.return %[[VAL]]
392+
ret float %1
393+
}
394+
395+
!0 = !{}
396+
397+
; // -----
398+
399+
; CHECK-LABEL: @invariant_group_store
400+
; CHECK-SAME: %[[VAL:[a-zA-Z0-9]+]]
401+
; CHECK-SAME: %[[PTR:[a-zA-Z0-9]+]]
402+
define void @invariant_group_store(float %val, ptr %ptr) {
403+
; CHECK: llvm.store %[[VAL]], %[[PTR]] invariant_group {alignment = 4 : i64} : f32, !llvm.ptr
404+
store float %val, ptr %ptr, align 4, !invariant.group !0
405+
; CHECK: llvm.return
406+
ret void
407+
}
408+
409+
!0 = !{}
410+
411+
; // -----
412+
386413
; CHECK-LABEL: @atomic_load_store
387414
; CHECK-SAME: %[[PTR:[a-zA-Z0-9]+]]
388415
define void @atomic_load_store(ptr %ptr) {

‎mlir/test/Target/LLVMIR/llvmir.mlir‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1967,6 +1967,20 @@ llvm.func @nontemporal_store_and_load(%ptr : !llvm.ptr) -> i32 {
19671967

19681968
// -----
19691969

1970+
// Check that invariant group attribute is exported as metadata node.
1971+
llvm.func @nontemporal_store_and_load(%ptr : !llvm.ptr) -> i32 {
1972+
%val = llvm.mlir.constant(42 : i32) : i32
1973+
// CHECK: store i32 42, ptr %{{.*}} !invariant.group ![[NODE:[0-9]+]]
1974+
llvm.store %val, %ptr invariant_group : i32, !llvm.ptr
1975+
// CHECK: %{{.*}} = load i32, ptr %{{.*}} !invariant.group ![[NODE]]
1976+
%1 = llvm.load %ptr invariant_group : !llvm.ptr -> i32
1977+
llvm.return %1 : i32
1978+
}
1979+
1980+
// CHECK: ![[NODE]] = !{}
1981+
1982+
// -----
1983+
19701984
llvm.func @atomic_store_and_load(%ptr : !llvm.ptr) {
19711985
// CHECK: load atomic
19721986
// CHECK-SAME: acquire, align 4

0 commit comments

Comments
 (0)