Skip to content

Commit 0019a9e

Browse files
tyb0807tyb0807
committed
[mlir][memref][vectpr] Add alignment attribute to memory access ops
1 parent 40d2f39 commit 0019a9e

File tree

6 files changed

+169
-6
lines changed

6 files changed

+169
-6
lines changed

‎mlir/docs/DefiningDialects/Operations.md

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -306,6 +306,8 @@ Right now, the following primitive constraints are supported:
306306
* `IntPositive`: Specifying an integer attribute whose value is positive
307307
* `IntNonNegative`: Specifying an integer attribute whose value is
308308
non-negative
309+
* `IntPowerOf2`: Specifying an integer attribute whose value is a power of
310+
two > 0
309311
* `ArrayMinCount<N>`: Specifying an array attribute to have at least `N`
310312
elements
311313
* `ArrayMaxCount<N>`: Specifying an array attribute to have at most `N`

‎mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td

Lines changed: 57 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1217,6 +1217,11 @@ def LoadOp : MemRef_Op<"load",
12171217
be reused in the cache. For details, refer to the
12181218
[https://llvm.org/docs/LangRef.html#load-instruction](LLVM load instruction).
12191219

1220+
An optional `alignment` attribute allows to specify the byte alignment of the
1221+
load operation. It must be a positive power of 2. The operation must access
1222+
memory at an address aligned to this boundary. Violations may lead to
1223+
architecture-specific faults or performance penalties.
1224+
A value of 0 indicates no specific alignment requirement.
12201225
Example:
12211226

12221227
```mlir
@@ -1227,7 +1232,39 @@ def LoadOp : MemRef_Op<"load",
12271232
let arguments = (ins Arg<AnyMemRef, "the reference to load from",
12281233
[MemRead]>:$memref,
12291234
Variadic<Index>:$indices,
1230-
DefaultValuedOptionalAttr<BoolAttr, "false">:$nontemporal);
1235+
DefaultValuedOptionalAttr<BoolAttr, "false">:$nontemporal,
1236+
ConfinedAttr<OptionalAttr<I64Attr>,
1237+
[AllAttrOf<[IntPositive, IntPowerOf2]>]>:$alignment);
1238+
1239+
let builders = [
1240+
OpBuilder<(ins "Value":$memref,
1241+
"ValueRange":$indices,
1242+
CArg<"bool", "false">:$nontemporal,
1243+
CArg<"uint64_t", "0">:$alignment), [{
1244+
return build($_builder, $_state, memref, indices, nontemporal,
1245+
alignment != 0 ? $_builder.getI64IntegerAttr(alignment) :
1246+
nullptr);
1247+
}]>,
1248+
OpBuilder<(ins "Type":$resultType,
1249+
"Value":$memref,
1250+
"ValueRange":$indices,
1251+
CArg<"bool", "false">:$nontemporal,
1252+
CArg<"uint64_t", "0">:$alignment), [{
1253+
return build($_builder, $_state, resultType, memref, indices, nontemporal,
1254+
alignment != 0 ? $_builder.getI64IntegerAttr(alignment) :
1255+
nullptr);
1256+
}]>,
1257+
OpBuilder<(ins "TypeRange":$resultTypes,
1258+
"Value":$memref,
1259+
"ValueRange":$indices,
1260+
CArg<"bool", "false">:$nontemporal,
1261+
CArg<"uint64_t", "0">:$alignment), [{
1262+
return build($_builder, $_state, resultTypes, memref, indices, nontemporal,
1263+
alignment != 0 ? $_builder.getI64IntegerAttr(alignment) :
1264+
nullptr);
1265+
}]>
1266+
];
1267+
12311268
let results = (outs AnyType:$result);
12321269

12331270
let extraClassDeclaration = [{
@@ -1913,6 +1950,11 @@ def MemRef_StoreOp : MemRef_Op<"store",
19131950
be reused in the cache. For details, refer to the
19141951
[https://llvm.org/docs/LangRef.html#store-instruction](LLVM store instruction).
19151952

1953+
An optional `alignment` attribute allows to specify the byte alignment of the
1954+
store operation. It must be a positive power of 2. The operation must access
1955+
memory at an address aligned to this boundary. Violations may lead to
1956+
architecture-specific faults or performance penalties.
1957+
A value of 0 indicates no specific alignment requirement.
19161958
Example:
19171959

19181960
```mlir
@@ -1924,13 +1966,25 @@ def MemRef_StoreOp : MemRef_Op<"store",
19241966
Arg<AnyMemRef, "the reference to store to",
19251967
[MemWrite]>:$memref,
19261968
Variadic<Index>:$indices,
1927-
DefaultValuedOptionalAttr<BoolAttr, "false">:$nontemporal);
1969+
DefaultValuedOptionalAttr<BoolAttr, "false">:$nontemporal,
1970+
ConfinedAttr<OptionalAttr<I64Attr>,
1971+
[AllAttrOf<[IntPositive, IntPowerOf2]>]>:$alignment);
19281972

19291973
let builders = [
1974+
OpBuilder<(ins "Value":$valueToStore,
1975+
"Value":$memref,
1976+
"ValueRange":$indices,
1977+
CArg<"bool", "false">:$nontemporal,
1978+
CArg<"uint64_t", "0">:$alignment), [{
1979+
return build($_builder, $_state, valueToStore, memref, indices, nontemporal,
1980+
alignment != 0 ? $_builder.getI64IntegerAttr(alignment) :
1981+
nullptr);
1982+
}]>,
19301983
OpBuilder<(ins "Value":$valueToStore, "Value":$memref), [{
19311984
$_state.addOperands(valueToStore);
19321985
$_state.addOperands(memref);
1933-
}]>];
1986+
}]>
1987+
];
19341988

19351989
let extraClassDeclaration = [{
19361990
Value getValueToStore() { return getOperand(0); }

‎mlir/include/mlir/Dialect/Vector/IR/VectorOps.td

Lines changed: 52 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1734,12 +1734,42 @@ def Vector_LoadOp : Vector_Op<"load"> {
17341734
```mlir
17351735
%result = vector.load %memref[%c0] : memref<7xf32>, vector<8xf32>
17361736
```
1737+
1738+
An optional `alignment` attribute allows to specify the byte alignment of the
1739+
load operation. It must be a positive power of 2. The operation must access
1740+
memory at an address aligned to this boundary. Violations may lead to
1741+
architecture-specific faults or performance penalties.
1742+
A value of 0 indicates no specific alignment requirement.
17371743
}];
17381744

17391745
let arguments = (ins Arg<AnyMemRef, "the reference to load from",
17401746
[MemRead]>:$base,
17411747
Variadic<Index>:$indices,
1742-
DefaultValuedOptionalAttr<BoolAttr, "false">:$nontemporal);
1748+
DefaultValuedOptionalAttr<BoolAttr, "false">:$nontemporal,
1749+
ConfinedAttr<OptionalAttr<I64Attr>,
1750+
[AllAttrOf<[IntPositive, IntPowerOf2]>]>:$alignment);
1751+
1752+
let builders = [
1753+
OpBuilder<(ins "VectorType":$resultType,
1754+
"Value":$base,
1755+
"ValueRange":$indices,
1756+
CArg<"bool", "false">:$nontemporal,
1757+
CArg<"uint64_t", "0">:$alignment), [{
1758+
return build($_builder, $_state, resultType, base, indices, nontemporal,
1759+
alignment != 0 ? $_builder.getI64IntegerAttr(alignment) :
1760+
nullptr);
1761+
}]>,
1762+
OpBuilder<(ins "TypeRange":$resultTypes,
1763+
"Value":$base,
1764+
"ValueRange":$indices,
1765+
CArg<"bool", "false">:$nontemporal,
1766+
CArg<"uint64_t", "0">:$alignment), [{
1767+
return build($_builder, $_state, resultTypes, base, indices, nontemporal,
1768+
alignment != 0 ? $_builder.getI64IntegerAttr(alignment) :
1769+
nullptr);
1770+
}]>
1771+
];
1772+
17431773
let results = (outs AnyVectorOfAnyRank:$result);
17441774

17451775
let extraClassDeclaration = [{
@@ -1818,15 +1848,34 @@ def Vector_StoreOp : Vector_Op<"store"> {
18181848
```mlir
18191849
vector.store %valueToStore, %memref[%c0] : memref<7xf32>, vector<8xf32>
18201850
```
1851+
1852+
An optional `alignment` attribute allows to specify the byte alignment of the
1853+
store operation. It must be a positive power of 2. The operation must access
1854+
memory at an address aligned to this boundary. Violations may lead to
1855+
architecture-specific faults or performance penalties.
1856+
A value of 0 indicates no specific alignment requirement.
18211857
}];
18221858

18231859
let arguments = (ins
18241860
AnyVectorOfAnyRank:$valueToStore,
18251861
Arg<AnyMemRef, "the reference to store to",
18261862
[MemWrite]>:$base,
18271863
Variadic<Index>:$indices,
1828-
DefaultValuedOptionalAttr<BoolAttr, "false">:$nontemporal
1829-
);
1864+
DefaultValuedOptionalAttr<BoolAttr, "false">:$nontemporal,
1865+
ConfinedAttr<OptionalAttr<I64Attr>,
1866+
[AllAttrOf<[IntPositive, IntPowerOf2]>]>:$alignment);
1867+
1868+
let builders = [
1869+
OpBuilder<(ins "Value":$valueToStore,
1870+
"Value":$base,
1871+
"ValueRange":$indices,
1872+
CArg<"bool", "false">:$nontemporal,
1873+
CArg<"uint64_t", "0">:$alignment), [{
1874+
return build($_builder, $_state, valueToStore, base, indices, nontemporal,
1875+
alignment != 0 ? $_builder.getI64IntegerAttr(alignment) :
1876+
nullptr);
1877+
}]>
1878+
];
18301879

18311880
let extraClassDeclaration = [{
18321881
MemRefType getMemRefType() {

‎mlir/include/mlir/IR/CommonAttrConstraints.td

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -796,6 +796,10 @@ def IntPositive : AttrConstraint<
796796
CPred<"::llvm::cast<::mlir::IntegerAttr>($_self).getValue().isStrictlyPositive()">,
797797
"whose value is positive">;
798798

799+
def IntPowerOf2 : AttrConstraint<
800+
CPred<"::llvm::cast<::mlir::IntegerAttr>($_self).getValue().isPowerOf2()">,
801+
"whose value is a power of two > 0">;
802+
799803
class ArrayMaxCount<int n> : AttrConstraint<
800804
CPred<"::llvm::cast<::mlir::ArrayAttr>($_self).size() <= " # n>,
801805
"with at most " # n # " elements">;

‎mlir/test/Dialect/MemRef/load-store-alignment.mlir

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,27 @@
1+
// RUN: mlir-opt -split-input-file -verify-diagnostics %s | FileCheck %s
2+
3+
// CHECK-LABEL: func @test_load_store_alignment
4+
// CHECK: memref.load {{.*}} {alignment = 16 : i64}
5+
// CHECK: memref.store {{.*}} {alignment = 16 : i64}
6+
func.func @test_load_store_alignment(%memref: memref<4xi32>) {
7+
%c0 = arith.constant 0 : index
8+
%val = memref.load %memref[%c0] { alignment = 16 } : memref<4xi32>
9+
memref.store %val, %memref[%c0] { alignment = 16 } : memref<4xi32>
10+
return
11+
}
12+
13+
// -----
14+
15+
func.func @test_invalid_negative_load_alignment(%memref: memref<4xi32>) {
16+
// expected-error @+1 {{custom op 'memref.load' 'memref.load' op attribute 'alignment' failed to satisfy constraint: 64-bit signless integer attribute whose value is positive and whose value is a power of two > 0}}
17+
%val = memref.load %memref[%c0] { alignment = -1 } : memref<4xi32>
18+
return
19+
}
20+
21+
// -----
22+
23+
func.func @test_invalid_non_power_of_2_store_alignment(%memref: memref<4xi32>, %val: i32) {
24+
// expected-error @+1 {{custom op 'memref.store' 'memref.store' op attribute 'alignment' failed to satisfy constraint: 64-bit signless integer attribute whose value is positive and whose value is a power of two > 0}}
25+
memref.store %val, %memref[%c0] { alignment = 1 } : memref<4xi32>
26+
return
27+
}

‎mlir/test/Dialect/Vector/load-store-alignment.mlir

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,27 @@
1+
// RUN: mlir-opt -split-input-file -verify-diagnostics %s | FileCheck %s
2+
3+
// CHECK-LABEL: func @test_load_store_alignment
4+
// CHECK: vector.load {{.*}} {alignment = 16 : i64}
5+
// CHECK: vector.store {{.*}} {alignment = 16 : i64}
6+
func.func @test_load_store_alignment(%memref: memref<4xi32>) {
7+
%c0 = arith.constant 0 : index
8+
%val = vector.load %memref[%c0] { alignment = 16 } : memref<4xi32>, vector<4xi32>
9+
vector.store %val, %memref[%c0] { alignment = 16 } : memref<4xi32>, vector<4xi32>
10+
return
11+
}
12+
13+
// -----
14+
15+
func.func @test_invalid_negative_load_alignment(%memref: memref<4xi32>) {
16+
// expected-error @+1 {{custom op 'vector.load' 'vector.load' op attribute 'alignment' failed to satisfy constraint: 64-bit signless integer attribute whose value is positive and whose value is a power of two > 0}}
17+
%val = vector.load %memref[%c0] { alignment = -1 } : memref<4xi32>, vector<4xi32>
18+
return
19+
}
20+
21+
// -----
22+
23+
func.func @test_invalid_non_power_of_2_store_alignment(%memref: memref<4xi32>, %val: vector<4xi32>) {
24+
// expected-error @+1 {{custom op 'vector.store' 'vector.store' op attribute 'alignment' failed to satisfy constraint: 64-bit signless integer attribute whose value is positive and whose value is a power of two > 0}}
25+
vector.store %val, %memref[%c0] { alignment = 1 } : memref<4xi32>, vector<4xi32>
26+
return
27+
}

0 commit comments

Comments
 (0)