Skip to content

Commit d361423

Browse files
committed
Add aligment attribute to vector.maskedload/store ops
1 parent 8a614df commit d361423

File tree

2 files changed

+66
-3
lines changed

2 files changed

+66
-3
lines changed

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

Lines changed: 51 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1899,7 +1899,9 @@ def Vector_MaskedLoadOp :
18991899
Arguments<(ins Arg<AnyMemRef, "", [MemRead]>:$base,
19001900
Variadic<Index>:$indices,
19011901
VectorOfNonZeroRankOf<[I1]>:$mask,
1902-
AnyVectorOfNonZeroRank:$pass_thru)>,
1902+
AnyVectorOfNonZeroRank:$pass_thru,
1903+
ConfinedAttr<OptionalAttr<I64Attr>,
1904+
[AllAttrOf<[IntPositive, IntPowerOf2]>]>:$alignment)>,
19031905
Results<(outs AnyVectorOfNonZeroRank:$result)> {
19041906

19051907
let summary = "loads elements from memory into a vector as defined by a mask vector";
@@ -1920,6 +1922,12 @@ def Vector_MaskedLoadOp :
19201922
comes from the pass-through vector regardless of the index, and the index is
19211923
allowed to be out-of-bounds.
19221924

1925+
An optional `alignment` attribute allows to specify the byte alignment of the
1926+
load operation. It must be a positive power of 2. The operation must access
1927+
memory at an address aligned to this boundary. Violations may lead to
1928+
architecture-specific faults or performance penalties.
1929+
A value of 0 indicates no specific alignment requirement.
1930+
19231931
The masked load can be used directly where applicable, or can be used
19241932
during progressively lowering to bring other memory operations closer to
19251933
hardware ISA support for a masked load. The semantics of the operation
@@ -1936,6 +1944,18 @@ def Vector_MaskedLoadOp :
19361944
: memref<?x?xf32>, vector<16xi1>, vector<16xf32> into vector<16xf32>
19371945
```
19381946
}];
1947+
let builders = [
1948+
OpBuilder<(ins "Type":$resultType,
1949+
"Value":$base,
1950+
"ValueRange":$indices,
1951+
"Value":$mask,
1952+
"Value":$pass_thru,
1953+
CArg<"uint64_t", "0">:$alignment), [{
1954+
return build($_builder, $_state, resultType, base, indices, mask, pass_thru,
1955+
alignment != 0 ? $_builder.getI64IntegerAttr(alignment) :
1956+
nullptr);
1957+
}]>
1958+
];
19391959
let extraClassDeclaration = [{
19401960
MemRefType getMemRefType() {
19411961
return ::llvm::cast<MemRefType>(getBase().getType());
@@ -1962,7 +1982,9 @@ def Vector_MaskedStoreOp :
19621982
Arguments<(ins Arg<AnyMemRef, "", [MemWrite]>:$base,
19631983
Variadic<Index>:$indices,
19641984
VectorOfNonZeroRankOf<[I1]>:$mask,
1965-
AnyVectorOfNonZeroRank:$valueToStore)> {
1985+
AnyVectorOfNonZeroRank:$valueToStore,
1986+
ConfinedAttr<OptionalAttr<I64Attr>,
1987+
[AllAttrOf<[IntPositive, IntPowerOf2]>]>:$alignment)> {
19661988

19671989
let summary = "stores elements from a vector into memory as defined by a mask vector";
19681990

@@ -1982,6 +2004,12 @@ def Vector_MaskedStoreOp :
19822004
is stored regardless of the index, and the index is allowed to be
19832005
out-of-bounds.
19842006

2007+
An optional `alignment` attribute allows to specify the byte alignment of the
2008+
store operation. It must be a positive power of 2. The operation must access
2009+
memory at an address aligned to this boundary. Violations may lead to
2010+
architecture-specific faults or performance penalties.
2011+
A value of 0 indicates no specific alignment requirement.
2012+
19852013
The masked store can be used directly where applicable, or can be used
19862014
during progressively lowering to bring other memory operations closer to
19872015
hardware ISA support for a masked store. The semantics of the operation
@@ -1998,6 +2026,27 @@ def Vector_MaskedStoreOp :
19982026
: memref<?x?xf32>, vector<16xi1>, vector<16xf32>
19992027
```
20002028
}];
2029+
let builders = [
2030+
OpBuilder<(ins "TypeRange":$resultTypes,
2031+
"Value":$base,
2032+
"ValueRange":$indices,
2033+
"Value":$mask,
2034+
"Value":$valueToStore,
2035+
CArg<"uint64_t", "0">:$alignment), [{
2036+
return build($_builder, $_state, resultTypes, base, indices, mask, valueToStore,
2037+
alignment != 0 ? $_builder.getI64IntegerAttr(alignment) :
2038+
nullptr);
2039+
}]>,
2040+
OpBuilder<(ins "Value":$base,
2041+
"ValueRange":$indices,
2042+
"Value":$mask,
2043+
"Value":$valueToStore,
2044+
CArg<"uint64_t", "0">:$alignment), [{
2045+
return build($_builder, $_state, base, indices, mask, valueToStore,
2046+
alignment != 0 ? $_builder.getI64IntegerAttr(alignment) :
2047+
nullptr);
2048+
}]>
2049+
];
20012050
let extraClassDeclaration = [{
20022051
MemRefType getMemRefType() {
20032052
return ::llvm::cast<MemRefType>(getBase().getType());

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

Lines changed: 15 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,17 @@
11
// RUN: mlir-opt -split-input-file -verify-diagnostics %s | FileCheck %s
22

3+
// CHECK-LABEL: func @test_masked_load_store_alignment
4+
// CHECK: vector.maskedload {{.*}} {alignment = 16 : i64}
5+
// CHECK: vector.maskedstore {{.*}} {alignment = 16 : i64}
6+
func.func @test_masked_load_store_alignment(%memref: memref<4xi32>, %mask: vector<4xi1>, %passthru: vector<4xi32>) {
7+
%c0 = arith.constant 0 : index
8+
%val = vector.maskedload %memref[%c0], %mask, %passthru { alignment = 16 } : memref<4xi32>, vector<4xi1>, vector<4xi32> into vector<4xi32>
9+
vector.maskedstore %memref[%c0], %mask, %val { alignment = 16 } : memref<4xi32>, vector<4xi1>, vector<4xi32>
10+
return
11+
}
12+
13+
// -----
14+
315
// CHECK-LABEL: func @test_load_store_alignment
416
// CHECK: vector.load {{.*}} {alignment = 16 : i64}
517
// CHECK: vector.store {{.*}} {alignment = 16 : i64}
@@ -13,6 +25,7 @@ func.func @test_load_store_alignment(%memref: memref<4xi32>) {
1325
// -----
1426

1527
func.func @test_invalid_negative_load_alignment(%memref: memref<4xi32>) {
28+
%c0 = arith.constant 0 : index
1629
// 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}}
1730
%val = vector.load %memref[%c0] { alignment = -1 } : memref<4xi32>, vector<4xi32>
1831
return
@@ -21,7 +34,8 @@ func.func @test_invalid_negative_load_alignment(%memref: memref<4xi32>) {
2134
// -----
2235

2336
func.func @test_invalid_non_power_of_2_store_alignment(%memref: memref<4xi32>, %val: vector<4xi32>) {
37+
%c0 = arith.constant 0 : index
2438
// 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>
39+
vector.store %val, %memref[%c0] { alignment = 3 } : memref<4xi32>, vector<4xi32>
2640
return
2741
}

0 commit comments

Comments
 (0)