Skip to content

Commit 77f2560

Browse files
fabianmcgfabianmcg
authored
[mlir][ptr] Add load and store ops. (#156093)
This patch adds the load and store operations to the ptr dialect. It's future work to implement SROA and Mem2Reg interfaces, as well as conversion to LLVM, and add alias information. This patch also fixes a bug in `OptionalProp` that was causing the bytecode writer to exit early of writing the Op props if an optional prop had the default value. Example: ```mlir func.func @load_ops(%arg0: !ptr.ptr<#ptr.generic_space>) -> (f32, f32, f32, f32, f32, i64, i32) { %0 = ptr.load %arg0 : !ptr.ptr<#ptr.generic_space> -> f32 %1 = ptr.load volatile %arg0 : !ptr.ptr<#ptr.generic_space> -> f32 %2 = ptr.load %arg0 nontemporal : !ptr.ptr<#ptr.generic_space> -> f32 %3 = ptr.load %arg0 invariant : !ptr.ptr<#ptr.generic_space> -> f32 %4 = ptr.load %arg0 invariant_group : !ptr.ptr<#ptr.generic_space> -> f32 %5 = ptr.load %arg0 atomic monotonic alignment = 8 : !ptr.ptr<#ptr.generic_space> -> i64 %6 = ptr.load volatile %arg0 atomic syncscope("workgroup") acquire nontemporal alignment = 4 : !ptr.ptr<#ptr.generic_space> -> i32 return %0, %1, %2, %3, %4, %5, %6 : f32, f32, f32, f32, f32, i64, i32 } func.func @store_ops(%arg0: !ptr.ptr<#ptr.generic_space>, %arg1: f32, %arg2: i64, %arg3: i32) { ptr.store %arg1, %arg0 : f32, !ptr.ptr<#ptr.generic_space> ptr.store volatile %arg1, %arg0 : f32, !ptr.ptr<#ptr.generic_space> ptr.store %arg1, %arg0 nontemporal : f32, !ptr.ptr<#ptr.generic_space> ptr.store %arg1, %arg0 invariant_group : f32, !ptr.ptr<#ptr.generic_space> ptr.store %arg2, %arg0 atomic monotonic alignment = 8 : i64, !ptr.ptr<#ptr.generic_space> ptr.store volatile %arg3, %arg0 atomic syncscope("workgroup") release nontemporal alignment = 4 : i32, !ptr.ptr<#ptr.generic_space> return } ``` Finally, this patch allows testing more advanced features of ptr memory spaces, for example: ```mlir // mlir-opt -verify-diagnostics func.func @store_const(%arg0: !ptr.ptr<#test.const_memory_space>, %arg1: i64) { // expected-error@+1 {{memory space is read-only}} ptr.store %arg1, %arg0 atomic monotonic alignment = 8 : i64, !ptr.ptr<#test.const_memory_space> return } ```
1 parent a3dfedf commit 77f2560

File tree

10 files changed

+341
-52
lines changed

10 files changed

+341
-52
lines changed

‎mlir/include/mlir/Dialect/Ptr/IR/MemorySpaceInterfaces.h‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -20,8 +20,8 @@
2020
namespace mlir {
2121
class Operation;
2222
namespace ptr {
23-
enum class AtomicBinOp : uint64_t;
24-
enum class AtomicOrdering : uint64_t;
23+
enum class AtomicBinOp : uint32_t;
24+
enum class AtomicOrdering : uint32_t;
2525
} // namespace ptr
2626
} // namespace mlir
2727

‎mlir/include/mlir/Dialect/Ptr/IR/MemorySpaceInterfaces.td‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -42,7 +42,7 @@ def MemorySpaceAttrInterface : AttrInterface<"MemorySpaceAttrInterface"> {
4242
/*methodName=*/ "isValidLoad",
4343
/*args=*/ (ins "::mlir::Type":$type,
4444
"::mlir::ptr::AtomicOrdering":$ordering,
45-
"::mlir::IntegerAttr":$alignment,
45+
"std::optional<int64_t>":$alignment,
4646
"::llvm::function_ref<::mlir::InFlightDiagnostic()>":$emitError)
4747
>,
4848
InterfaceMethod<
@@ -57,7 +57,7 @@ def MemorySpaceAttrInterface : AttrInterface<"MemorySpaceAttrInterface"> {
5757
/*methodName=*/ "isValidStore",
5858
/*args=*/ (ins "::mlir::Type":$type,
5959
"::mlir::ptr::AtomicOrdering":$ordering,
60-
"::mlir::IntegerAttr":$alignment,
60+
"std::optional<int64_t>":$alignment,
6161
"::llvm::function_ref<::mlir::InFlightDiagnostic()>":$emitError)
6262
>,
6363
InterfaceMethod<
@@ -73,7 +73,7 @@ def MemorySpaceAttrInterface : AttrInterface<"MemorySpaceAttrInterface"> {
7373
/*args=*/ (ins "::mlir::ptr::AtomicBinOp":$op,
7474
"::mlir::Type":$type,
7575
"::mlir::ptr::AtomicOrdering":$ordering,
76-
"::mlir::IntegerAttr":$alignment,
76+
"std::optional<int64_t>":$alignment,
7777
"::llvm::function_ref<::mlir::InFlightDiagnostic()>":$emitError)
7878
>,
7979
InterfaceMethod<
@@ -90,7 +90,7 @@ def MemorySpaceAttrInterface : AttrInterface<"MemorySpaceAttrInterface"> {
9090
/*args=*/ (ins "::mlir::Type":$type,
9191
"::mlir::ptr::AtomicOrdering":$successOrdering,
9292
"::mlir::ptr::AtomicOrdering":$failureOrdering,
93-
"::mlir::IntegerAttr":$alignment,
93+
"std::optional<int64_t>":$alignment,
9494
"::llvm::function_ref<::mlir::InFlightDiagnostic()>":$emitError)
9595
>,
9696
InterfaceMethod<

‎mlir/include/mlir/Dialect/Ptr/IR/PtrEnums.td‎

Lines changed: 28 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -15,25 +15,25 @@ include "mlir/IR/EnumAttr.td"
1515
// Atomic binary op enum attribute.
1616
//===----------------------------------------------------------------------===//
1717

18-
def AtomicBinOpXchg : I64EnumAttrCase<"xchg", 0, "xchg">;
19-
def AtomicBinOpAdd : I64EnumAttrCase<"add", 1, "add">;
20-
def AtomicBinOpSub : I64EnumAttrCase<"sub", 2, "sub">;
21-
def AtomicBinOpAnd : I64EnumAttrCase<"_and", 3, "_and">;
22-
def AtomicBinOpNand : I64EnumAttrCase<"nand", 4, "nand">;
23-
def AtomicBinOpOr : I64EnumAttrCase<"_or", 5, "_or">;
24-
def AtomicBinOpXor : I64EnumAttrCase<"_xor", 6, "_xor">;
25-
def AtomicBinOpMax : I64EnumAttrCase<"max", 7, "max">;
26-
def AtomicBinOpMin : I64EnumAttrCase<"min", 8, "min">;
27-
def AtomicBinOpUMax : I64EnumAttrCase<"umax", 9, "umax">;
28-
def AtomicBinOpUMin : I64EnumAttrCase<"umin", 10, "umin">;
29-
def AtomicBinOpFAdd : I64EnumAttrCase<"fadd", 11, "fadd">;
30-
def AtomicBinOpFSub : I64EnumAttrCase<"fsub", 12, "fsub">;
31-
def AtomicBinOpFMax : I64EnumAttrCase<"fmax", 13, "fmax">;
32-
def AtomicBinOpFMin : I64EnumAttrCase<"fmin", 14, "fmin">;
33-
def AtomicBinOpUIncWrap : I64EnumAttrCase<"uinc_wrap", 15, "uinc_wrap">;
34-
def AtomicBinOpUDecWrap : I64EnumAttrCase<"udec_wrap", 16, "udec_wrap">;
18+
def AtomicBinOpXchg : I32EnumCase<"xchg", 0, "xchg">;
19+
def AtomicBinOpAdd : I32EnumCase<"add", 1, "add">;
20+
def AtomicBinOpSub : I32EnumCase<"sub", 2, "sub">;
21+
def AtomicBinOpAnd : I32EnumCase<"_and", 3, "_and">;
22+
def AtomicBinOpNand : I32EnumCase<"nand", 4, "nand">;
23+
def AtomicBinOpOr : I32EnumCase<"_or", 5, "_or">;
24+
def AtomicBinOpXor : I32EnumCase<"_xor", 6, "_xor">;
25+
def AtomicBinOpMax : I32EnumCase<"max", 7, "max">;
26+
def AtomicBinOpMin : I32EnumCase<"min", 8, "min">;
27+
def AtomicBinOpUMax : I32EnumCase<"umax", 9, "umax">;
28+
def AtomicBinOpUMin : I32EnumCase<"umin", 10, "umin">;
29+
def AtomicBinOpFAdd : I32EnumCase<"fadd", 11, "fadd">;
30+
def AtomicBinOpFSub : I32EnumCase<"fsub", 12, "fsub">;
31+
def AtomicBinOpFMax : I32EnumCase<"fmax", 13, "fmax">;
32+
def AtomicBinOpFMin : I32EnumCase<"fmin", 14, "fmin">;
33+
def AtomicBinOpUIncWrap : I32EnumCase<"uinc_wrap", 15, "uinc_wrap">;
34+
def AtomicBinOpUDecWrap : I32EnumCase<"udec_wrap", 16, "udec_wrap">;
3535

36-
def AtomicBinOp : I64EnumAttr<
36+
def AtomicBinOp : I32Enum<
3737
"AtomicBinOp",
3838
"ptr.atomicrmw binary operations",
3939
[AtomicBinOpXchg, AtomicBinOpAdd, AtomicBinOpSub, AtomicBinOpAnd,
@@ -48,15 +48,15 @@ def AtomicBinOp : I64EnumAttr<
4848
// Atomic ordering enum attribute.
4949
//===----------------------------------------------------------------------===//
5050

51-
def AtomicOrderingNotAtomic : I64EnumAttrCase<"not_atomic", 0, "not_atomic">;
52-
def AtomicOrderingUnordered : I64EnumAttrCase<"unordered", 1, "unordered">;
53-
def AtomicOrderingMonotonic : I64EnumAttrCase<"monotonic", 2, "monotonic">;
54-
def AtomicOrderingAcquire : I64EnumAttrCase<"acquire", 3, "acquire">;
55-
def AtomicOrderingRelease : I64EnumAttrCase<"release", 4, "release">;
56-
def AtomicOrderingAcqRel : I64EnumAttrCase<"acq_rel", 5, "acq_rel">;
57-
def AtomicOrderingSeqCst : I64EnumAttrCase<"seq_cst", 6, "seq_cst">;
51+
def AtomicOrderingNotAtomic : I32EnumCase<"not_atomic", 0, "not_atomic">;
52+
def AtomicOrderingUnordered : I32EnumCase<"unordered", 1, "unordered">;
53+
def AtomicOrderingMonotonic : I32EnumCase<"monotonic", 2, "monotonic">;
54+
def AtomicOrderingAcquire : I32EnumCase<"acquire", 3, "acquire">;
55+
def AtomicOrderingRelease : I32EnumCase<"release", 4, "release">;
56+
def AtomicOrderingAcqRel : I32EnumCase<"acq_rel", 5, "acq_rel">;
57+
def AtomicOrderingSeqCst : I32EnumCase<"seq_cst", 6, "seq_cst">;
5858

59-
def AtomicOrdering : I64EnumAttr<
59+
def AtomicOrdering : I32Enum<
6060
"AtomicOrdering",
6161
"Atomic ordering for LLVM's memory model",
6262
[AtomicOrderingNotAtomic, AtomicOrderingUnordered, AtomicOrderingMonotonic,
@@ -66,6 +66,8 @@ def AtomicOrdering : I64EnumAttr<
6666
let cppNamespace = "::mlir::ptr";
6767
}
6868

69+
def AtomicOrderingProp : EnumProp<AtomicOrdering>;
70+
6971
//===----------------------------------------------------------------------===//
7072
// Ptr add flags enum properties.
7173
//===----------------------------------------------------------------------===//

‎mlir/include/mlir/Dialect/Ptr/IR/PtrOps.td‎

Lines changed: 127 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -118,6 +118,133 @@ def Ptr_PtrAddOp : Pointer_Op<"ptr_add", [
118118
}];
119119
}
120120

121+
//===----------------------------------------------------------------------===//
122+
// LoadOp
123+
//===----------------------------------------------------------------------===//
124+
125+
def AlignmentProp : OptionalProp<I64Prop>;
126+
127+
def Ptr_LoadOp : Pointer_Op<"load", [
128+
DeclareOpInterfaceMethods<MemoryEffectsOpInterface>
129+
]> {
130+
let description = [{
131+
The `load` operation is used to read from memory. A load may be marked as
132+
atomic, volatile, and/or nontemporal.
133+
134+
An atomic load only supports a limited set of value types, and requires
135+
an explicit alignment.
136+
137+
Examples:
138+
```mlir
139+
// A volatile load of a float variable.
140+
%0 = ptr.load volatile %ptr : !ptr.ptr -> f32
141+
142+
// A nontemporal load of a float variable.
143+
%0 = ptr.load %ptr nontemporal : !ptr.ptr -> f32
144+
145+
// An atomic load of an integer variable.
146+
%0 = ptr.load %ptr atomic monotonic alignment = 8 : !ptr.ptr -> i64
147+
```
148+
149+
See the following link for more details on the meaning of `alignment`,
150+
`volatile_`, `nontemporal`, `invariant`, `invariant_group`, `ordering`,
151+
and `syncscope`:
152+
https://llvm.org/docs/LangRef.html#load-instruction
153+
}];
154+
let arguments = (ins Ptr_PtrType:$ptr,
155+
AlignmentProp:$alignment,
156+
UnitProp:$volatile_,
157+
UnitProp:$nontemporal,
158+
UnitProp:$invariant,
159+
UnitProp:$invariantGroup,
160+
DefaultValuedProp<
161+
AtomicOrderingProp,
162+
"AtomicOrdering::not_atomic">:$ordering,
163+
OptionalAttr<StrAttr>:$syncscope);
164+
let results = (outs AnyType:$value);
165+
let assemblyFormat = [{
166+
(`volatile` $volatile_^)? $ptr
167+
(`atomic` (`syncscope` `(` $syncscope^ `)`)? $ordering^)?
168+
oilist(
169+
`nontemporal` $nontemporal |
170+
`invariant` $invariant |
171+
`invariant_group` $invariantGroup |
172+
`alignment` `=` $alignment
173+
)
174+
attr-dict `:` qualified(type($ptr)) `->` type($value)
175+
}];
176+
let builders = [
177+
OpBuilder<(ins "Type":$type, "Value":$ptr,
178+
CArg<"unsigned", "0">:$alignment, CArg<"bool", "false">:$isVolatile,
179+
CArg<"bool", "false">:$isNonTemporal, CArg<"bool", "false">:$isInvariant,
180+
CArg<"bool", "false">:$isInvariantGroup,
181+
CArg<"AtomicOrdering", "AtomicOrdering::not_atomic">:$ordering,
182+
CArg<"StringRef", "StringRef()">:$syncscope)>
183+
];
184+
let hasVerifier = 1;
185+
}
186+
187+
//===----------------------------------------------------------------------===//
188+
// StoreOp
189+
//===----------------------------------------------------------------------===//
190+
191+
def Ptr_StoreOp : Pointer_Op<"store", [
192+
DeclareOpInterfaceMethods<MemoryEffectsOpInterface>
193+
]> {
194+
let description = [{
195+
The `store` operation is used to write to memory. A store may be marked as
196+
atomic, volatile, and/or nontemporal.
197+
198+
An atomic store only supports a limited set of value types, and requires
199+
an explicit alignment.
200+
201+
Examples:
202+
```mlir
203+
// A volatile store of a float variable.
204+
ptr.store volatile %val, %ptr : f32, !ptr.ptr
205+
206+
// A nontemporal store of a float variable.
207+
ptr.store %val, %ptr nontemporal : f32, !ptr.ptr
208+
209+
// An atomic store of an integer variable.
210+
ptr.store %val, %ptr atomic monotonic alignment = 8: i64, !ptr.ptr
211+
```
212+
213+
See the following link for more details on the meaning of `alignment`,
214+
`volatile_`, `nontemporal`, `invariant_group`, `ordering`, and `syncscope`:
215+
https://llvm.org/docs/LangRef.html#store-instruction
216+
}];
217+
let arguments = (ins AnyType:$value,
218+
Ptr_PtrType:$ptr,
219+
AlignmentProp:$alignment,
220+
UnitProp:$volatile_,
221+
UnitProp:$nontemporal,
222+
UnitProp:$invariantGroup,
223+
DefaultValuedProp<
224+
AtomicOrderingProp,
225+
"AtomicOrdering::not_atomic">:$ordering,
226+
OptionalAttr<StrAttr>:$syncscope);
227+
let assemblyFormat = [{
228+
(`volatile` $volatile_^)? $value `,` $ptr
229+
(`atomic` (`syncscope` `(` $syncscope^ `)`)? $ordering^)?
230+
oilist(
231+
`nontemporal` $nontemporal |
232+
`invariant_group` $invariantGroup |
233+
`alignment` `=` $alignment
234+
)
235+
attr-dict `:` type($value) `,` qualified(type($ptr))
236+
}];
237+
let builders = [
238+
OpBuilder<(ins "Value":$value, "Value":$ptr,
239+
CArg<"unsigned", "0">:$alignment, CArg<"bool", "false">:$isVolatile,
240+
CArg<"bool", "false">:$isNonTemporal,
241+
CArg<"bool", "false">:$isInvariantGroup,
242+
CArg<"AtomicOrdering", "AtomicOrdering::not_atomic">:$ordering,
243+
CArg<"StringRef", "StringRef()">:$syncscope)>
244+
];
245+
let hasVerifier = 1;
246+
}
247+
121248
//===----------------------------------------------------------------------===//
122249
// ToPtrOp
123250
//===----------------------------------------------------------------------===//

‎mlir/include/mlir/IR/Properties.td‎

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -773,9 +773,10 @@ class OptionalProp<Property p, bit canDelegateParsing = 1>
773773
}];
774774
let writeToMlirBytecode = [{
775775
$_writer.writeOwnedBool($_storage.has_value());
776-
if (!$_storage.has_value())
777-
return;
778-
}] # !subst("$_storage", "(*($_storage))", p.writeToMlirBytecode);
776+
if ($_storage.has_value()) {
777+
}] # !subst("$_storage", "(*($_storage))", p.writeToMlirBytecode) # [{
778+
}
779+
}];
779780

780781
let hashProperty = !if(!empty(p.hashProperty), p.hashProperty,
781782
[{ hash_value($_storage.has_value() ? std::optional<::llvm::hash_code>{}] #

‎mlir/lib/Dialect/Ptr/IR/PtrAttrs.cpp‎

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -22,26 +22,27 @@ constexpr const static unsigned kBitsInByte = 8;
2222
//===----------------------------------------------------------------------===//
2323

2424
bool GenericSpaceAttr::isValidLoad(
25-
Type type, ptr::AtomicOrdering ordering, IntegerAttr alignment,
25+
Type type, ptr::AtomicOrdering ordering, std::optional<int64_t> alignment,
2626
function_ref<InFlightDiagnostic()> emitError) const {
2727
return true;
2828
}
2929

3030
bool GenericSpaceAttr::isValidStore(
31-
Type type, ptr::AtomicOrdering ordering, IntegerAttr alignment,
31+
Type type, ptr::AtomicOrdering ordering, std::optional<int64_t> alignment,
3232
function_ref<InFlightDiagnostic()> emitError) const {
3333
return true;
3434
}
3535

3636
bool GenericSpaceAttr::isValidAtomicOp(
3737
ptr::AtomicBinOp op, Type type, ptr::AtomicOrdering ordering,
38-
IntegerAttr alignment, function_ref<InFlightDiagnostic()> emitError) const {
38+
std::optional<int64_t> alignment,
39+
function_ref<InFlightDiagnostic()> emitError) const {
3940
return true;
4041
}
4142

4243
bool GenericSpaceAttr::isValidAtomicXchg(
4344
Type type, ptr::AtomicOrdering successOrdering,
44-
ptr::AtomicOrdering failureOrdering, IntegerAttr alignment,
45+
ptr::AtomicOrdering failureOrdering, std::optional<int64_t> alignment,
4546
function_ref<InFlightDiagnostic()> emitError) const {
4647
return true;
4748
}

0 commit comments

Comments
 (0)