Skip to content

[mlir][vector] Add alignment attribute to vector operations. #152507

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 17 commits into from
Aug 25, 2025
Merged

[mlir][vector] Add alignment attribute to vector operations. #152507

Show file tree
Hide file tree
Changes from 11 commits
Commits
Show all changes
17 commits
Select commit Hold shift + click to select a range
a97cbc6
[mlir][vector] Add alignment to vector.gather.
amd-eochoalo Aug 7, 2025
158b91c
[mlir][vector] Add alignment to vector.scatter
amd-eochoalo Aug 7, 2025
36949d1
[mlir][vector] Add alignment to compressstore
amd-eochoalo Aug 7, 2025
a4d820f
[mlir][vector] Add alignment to expandload
amd-eochoalo Aug 7, 2025
e2ad0f9
Use llvm::MaybeAlign
amd-eochoalo Aug 8, 2025
92b3886
Add documentation for alignment attribute
amd-eochoalo Aug 9, 2025
482ad75
Change names in test functions
amd-eochoalo Aug 11, 2025
9b451db
Remove unnecessary constructor
amd-eochoalo Aug 11, 2025
3180cd0
Remove documentation from wrong op
amd-eochoalo Aug 11, 2025
47db5b1
Fix documention
amd-eochoalo Aug 11, 2025
a714521
Remove unnecessary constructor
amd-eochoalo Aug 11, 2025
0f81b89
Update mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
amd-eochoalo Aug 15, 2025
be906aa
Merge branch 'main' into eochoa/2025-08-07/vector-ops-alignment-minus…
amd-eochoalo Aug 15, 2025
46d9010
Remove 'value of 0 indicates no alignment requirement.'
amd-eochoalo Aug 15, 2025
87958eb
Change tests from negative to non-positive
amd-eochoalo Aug 15, 2025
c80adf0
Merge branch 'main' into eochoa/2025-08-06/vector-ops-alignment
amd-eochoalo Aug 20, 2025
dcf38e2
Merge branch 'main' into eochoa/2025-08-06/vector-ops-alignment
amd-eochoalo Aug 25, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
91 changes: 87 additions & 4 deletions mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2054,7 +2054,9 @@ def Vector_GatherOp :
Variadic<Index>:$indices,
VectorOfNonZeroRankOf<[AnyInteger, Index]>:$index_vec,
VectorOfNonZeroRankOf<[I1]>:$mask,
AnyVectorOfNonZeroRank:$pass_thru)>,
AnyVectorOfNonZeroRank:$pass_thru,
ConfinedAttr<OptionalAttr<I64Attr>,
[AllAttrOf<[IntPositive, IntPowerOf2]>]>:$alignment)>,
Results<(outs AnyVectorOfNonZeroRank:$result)> {

let summary = [{
Expand Down Expand Up @@ -2085,6 +2087,12 @@ def Vector_GatherOp :
during progressively lowering to bring other memory operations closer to
hardware ISA support for a gather.

An optional `alignment` attribute allows to specify the byte alignment of the
scatter operation. It must be a positive power of 2. The operation must access
memory at an address aligned to this boundary. Violations may lead to
architecture-specific faults or performance penalties.
A value of 0 indicates no specific alignment requirement.

Examples:

```mlir
Expand All @@ -2111,6 +2119,20 @@ def Vector_GatherOp :
"`into` type($result)";
let hasCanonicalizer = 1;
let hasVerifier = 1;

let builders = [
OpBuilder<(ins "VectorType":$resultType,
"Value":$base,
"ValueRange":$indices,
"Value":$index_vec,
"Value":$mask,
"Value":$passthrough,
CArg<"llvm::MaybeAlign", "llvm::MaybeAlign()">:$alignment), [{
return build($_builder, $_state, resultType, base, indices, index_vec, mask, passthrough,
alignment.has_value() ? $_builder.getI64IntegerAttr(alignment->value()) :
nullptr);
}]>
];
}

def Vector_ScatterOp :
Expand All @@ -2119,7 +2141,9 @@ def Vector_ScatterOp :
Variadic<Index>:$indices,
VectorOfNonZeroRankOf<[AnyInteger, Index]>:$index_vec,
VectorOfNonZeroRankOf<[I1]>:$mask,
AnyVectorOfNonZeroRank:$valueToStore)> {
AnyVectorOfNonZeroRank:$valueToStore,
ConfinedAttr<OptionalAttr<I64Attr>,
[AllAttrOf<[IntPositive, IntPowerOf2]>]>:$alignment)> {

let summary = [{
scatters elements from a vector into memory as defined by an index vector
Expand Down Expand Up @@ -2153,6 +2177,12 @@ def Vector_ScatterOp :
correspond to those of the `llvm.masked.scatter`
[intrinsic](https://llvm.org/docs/LangRef.html#llvm-masked-scatter-intrinsics).

An optional `alignment` attribute allows to specify the byte alignment of the
scatter operation. It must be a positive power of 2. The operation must access
memory at an address aligned to this boundary. Violations may lead to
architecture-specific faults or performance penalties.
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The violation is strangely defined here, why isn't this specified as UB?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

In this particular case, there we some operations defined in this other PR #144344 which used this wording as documentation. I think changing the wording to be undefined behaviour is reasonable. Would something like the following be preferrable?

    Violations will result in undefined behaviour and may lead to
    architecture-specific faults or performance penalties.

Copy link
Collaborator

@joker-eph joker-eph Aug 25, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think

Violating this requirements triggers immediate undefined behavior
seems reasonable for the stores.

For the loads maybe we should instead use:

Violating this requirements will make the loaded value a poison value.

or something like that.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think immediate UB is more suitable since using aligned loads on unaligned pointers is known to crash on some architectures

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good point, but that means we can't speculate an aligned load anymore right?
Maybe that's expected and OK though.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I can open a PR to change the wording. Thanks @joker-eph !

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I checked LLVM and that seems like what we expect:

It is the responsibility of the code emitter to ensure that the alignment information is correct. Overestimating the alignment results in undefined behavior. Underestimating the alignment may produce less efficient code. An alignment of 1 is always safe.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

may lead to architecture-specific faults or performance penalties

I personally prefer to avoid phrases like this unless we can provide a specific example. More generic terms like "UB" or "poison" tend to be more universal.

A value of 0 indicates no specific alignment requirement.

Examples:

```mlir
Expand All @@ -2177,14 +2207,29 @@ def Vector_ScatterOp :
"type($index_vec) `,` type($mask) `,` type($valueToStore)";
let hasCanonicalizer = 1;
let hasVerifier = 1;

let builders = [
OpBuilder<(ins "Value":$base,
"ValueRange":$indices,
"Value":$index_vec,
"Value":$mask,
"Value":$valueToStore,
CArg<"llvm::MaybeAlign", "llvm::MaybeAlign()">: $alignment), [{
return build($_builder, $_state, base, indices, index_vec, mask, valueToStore,
alignment.has_value() ? $_builder.getI64IntegerAttr(alignment->value()) :
nullptr);
}]>
];
}

def Vector_ExpandLoadOp :
Vector_Op<"expandload">,
Arguments<(ins Arg<AnyMemRef, "", [MemRead]>:$base,
Variadic<Index>:$indices,
FixedVectorOfNonZeroRankOf<[I1]>:$mask,
AnyVectorOfNonZeroRank:$pass_thru)>,
AnyVectorOfNonZeroRank:$pass_thru,
ConfinedAttr<OptionalAttr<I64Attr>,
[AllAttrOf<[IntPositive, IntPowerOf2]>]>:$alignment)>,
Results<(outs AnyVectorOfNonZeroRank:$result)> {

let summary = "reads elements from memory and spreads them into a vector as defined by a mask";
Expand Down Expand Up @@ -2216,6 +2261,12 @@ def Vector_ExpandLoadOp :
correspond to those of the `llvm.masked.expandload`
[intrinsic](https://llvm.org/docs/LangRef.html#llvm-masked-expandload-intrinsics).

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

Note, at the moment this Op is only available for fixed-width vectors.

Examples:
Expand Down Expand Up @@ -2246,14 +2297,29 @@ def Vector_ExpandLoadOp :
"type($base) `,` type($mask) `,` type($pass_thru) `into` type($result)";
let hasCanonicalizer = 1;
let hasVerifier = 1;

let builders = [
OpBuilder<(ins "VectorType":$resultType,
"Value":$base,
"ValueRange":$indices,
"Value":$mask,
"Value":$passthrough,
CArg<"llvm::MaybeAlign", "llvm::MaybeAlign()">:$alignment), [{
return build($_builder, $_state, resultType, base, indices, mask, passthrough,
alignment.has_value() ? $_builder.getI64IntegerAttr(alignment->value()) :
nullptr);
}]>
];
}

def Vector_CompressStoreOp :
Vector_Op<"compressstore">,
Arguments<(ins Arg<AnyMemRef, "", [MemWrite]>:$base,
Variadic<Index>:$indices,
FixedVectorOfNonZeroRankOf<[I1]>:$mask,
AnyVectorOfNonZeroRank:$valueToStore)> {
AnyVectorOfNonZeroRank:$valueToStore,
ConfinedAttr<OptionalAttr<I64Attr>,
[AllAttrOf<[IntPositive, IntPowerOf2]>]>:$alignment)> {

let summary = "writes elements selectively from a vector as defined by a mask";

Expand Down Expand Up @@ -2284,6 +2350,12 @@ def Vector_CompressStoreOp :
correspond to those of the `llvm.masked.compressstore`
[intrinsic](https://llvm.org/docs/LangRef.html#llvm-masked-compressstore-intrinsics).

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

Note, at the moment this Op is only available for fixed-width vectors.

Examples:
Expand Down Expand Up @@ -2312,6 +2384,17 @@ def Vector_CompressStoreOp :
"type($base) `,` type($mask) `,` type($valueToStore)";
let hasCanonicalizer = 1;
let hasVerifier = 1;
let builders = [
OpBuilder<(ins "Value":$base,
"ValueRange":$indices,
"Value":$mask,
"Value":$valueToStore,
CArg<"llvm::MaybeAlign", "llvm::MaybeAlign()">:$alignment), [{
return build($_builder, $_state, base, indices, valueToStore, mask,
alignment.has_value() ? $_builder.getI64IntegerAttr(alignment->value()) :
nullptr);
}]>
];
}

def Vector_ShapeCastOp :
Expand Down
76 changes: 70 additions & 6 deletions mlir/test/Dialect/Vector/invalid.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1317,7 +1317,7 @@ func.func @maskedload_negative_alignment(%base: memref<4xi32>, %mask: vector<32x

// -----

func.func @maskedload_nonpoweroftwo_alignment(%base: memref<4xi32>, %mask: vector<32xi1>, %pass: vector<1xi32>, %index: index) {
func.func @maskedload_non_power_of_2_alignment(%base: memref<4xi32>, %mask: vector<32xi1>, %pass: vector<1xi32>, %index: index) {
// expected-error@below {{'vector.maskedload' 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}}
%val = vector.maskedload %base[%index], %mask, %pass { alignment = 3 } : memref<4xi32>, vector<32xi1>, vector<1xi32> into vector<1xi32>
return
Expand Down Expand Up @@ -1368,7 +1368,7 @@ func.func @maskedstore_negative_alignment(%base: memref<4xi32>, %mask: vector<32

// -----

func.func @maskedstore_nonpoweroftwo_alignment(%base: memref<4xi32>, %mask: vector<32xi1>, %value: vector<1xi32>, %index: index) {
func.func @maskedstore_non_power_of_2_alignment(%base: memref<4xi32>, %mask: vector<32xi1>, %value: vector<1xi32>, %index: index) {
// expected-error@below {{'vector.maskedstore' 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}}
vector.maskedstore %base[%index], %mask, %value { alignment = 3 } : memref<4xi32>, vector<32xi1>, vector<1xi32> into vector<1xi32>
return
Expand Down Expand Up @@ -1470,6 +1470,24 @@ func.func @gather_pass_thru_type_mismatch(%base: memref<?xf32>, %indices: vector

// -----

func.func @gather_negative_alignment(%base: memref<16xf32>, %indices: vector<16xi32>,
%mask: vector<16xi1>, %pass_thru: vector<16xf32>, %c0 : index) {
// expected-error@+2 {{'vector.gather' 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}}
%0 = vector.gather %base[%c0][%indices], %mask, %pass_thru
{ alignment = -1 } : memref<16xf32>, vector<16xi32>, vector<16xi1>, vector<16xf32> into vector<16xf32>
}

// -----

func.func @gather_non_power_of_two_alignment(%base: memref<16xf32>, %indices: vector<16xi32>,
%mask: vector<16xi1>, %pass_thru: vector<16xf32>, %c0 : index) {
// expected-error@+2 {{'vector.gather' 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}}
%0 = vector.gather %base[%c0][%indices], %mask, %pass_thru
{ alignment = 3 } : memref<16xf32>, vector<16xi32>, vector<16xi1>, vector<16xf32> into vector<16xf32>
}

// -----

func.func @scatter_to_vector(%base: vector<16xf32>, %indices: vector<16xi32>,
%mask: vector<16xi1>, %pass_thru: vector<16xf32>) {
%c0 = arith.constant 0 : index
Expand Down Expand Up @@ -1531,6 +1549,24 @@ func.func @scatter_dim_mask_mismatch(%base: memref<?xf32>, %indices: vector<16xi

// -----

func.func @scatter_negative_alignment(%base: memref<?xf32>, %indices: vector<16xi32>,
%mask: vector<16xi1>, %value: vector<16xf32>, %c0: index) {
// expected-error@+1 {{'vector.scatter' 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}}
vector.scatter %base[%c0][%indices], %mask, %value { alignment = -1 }
: memref<?xf32>, vector<16xi32>, vector<16xi1>, vector<16xf32>
}

// -----

func.func @scatter_non_power_of_2_alignment(%base: memref<?xf32>, %indices: vector<16xi32>,
%mask: vector<16xi1>, %value: vector<16xf32>, %c0: index) {
// expected-error@+1 {{'vector.scatter' 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}}
vector.scatter %base[%c0][%indices], %mask, %value { alignment = 3 }
: memref<?xf32>, vector<16xi32>, vector<16xi1>, vector<16xf32>
}

// -----

func.func @expand_base_type_mismatch(%base: memref<?xf64>, %mask: vector<16xi1>, %pass_thru: vector<16xf32>) {
%c0 = arith.constant 0 : index
// expected-error@+1 {{'vector.expandload' op base and result element type should match}}
Expand Down Expand Up @@ -1571,6 +1607,20 @@ func.func @expand_memref_mismatch(%base: memref<?x?xf32>, %mask: vector<16xi1>,

// -----

func.func @expand_negative_alignment(%base: memref<?xf32>, %mask: vector<16xi1>, %pass_thru: vector<16xf32>, %c0: index) {
// expected-error@+1 {{'vector.expandload' 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}}
%0 = vector.expandload %base[%c0], %mask, %pass_thru { alignment = -1 } : memref<?xf32>, vector<16xi1>, vector<16xf32> into vector<16xf32>
}

// -----

func.func @expand_non_power_of_2_alignment(%base: memref<?xf32>, %mask: vector<16xi1>, %pass_thru: vector<16xf32>, %c0: index) {
// expected-error@+1 {{'vector.expandload' 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}}
%0 = vector.expandload %base[%c0], %mask, %pass_thru { alignment = 3 } : memref<?xf32>, vector<16xi1>, vector<16xf32> into vector<16xf32>
}

// -----

func.func @compress_base_type_mismatch(%base: memref<?xf64>, %mask: vector<16xi1>, %value: vector<16xf32>) {
%c0 = arith.constant 0 : index
// expected-error@+1 {{'vector.compressstore' op base and valueToStore element type should match}}
Expand Down Expand Up @@ -1603,6 +1653,20 @@ func.func @compress_memref_mismatch(%base: memref<?x?xf32>, %mask: vector<16xi1>

// -----

func.func @compress_negative_alignment(%base: memref<?xf32>, %mask: vector<16xi1>, %value: vector<16xf32>, %c0: index) {
// expected-error @below {{'vector.compressstore' 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}}
vector.compressstore %base[%c0], %mask, %value { alignment = -1 } : memref<?xf32>, vector<16xi1>, vector<16xf32>
}

// -----

func.func @compress_non_power_of_2_alignment(%base: memref<?xf32>, %mask: vector<16xi1>, %value: vector<16xf32>, %c0: index) {
// expected-error @below {{'vector.compressstore' 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}}
vector.compressstore %base[%c0], %mask, %value { alignment = 3 } : memref<?xf32>, vector<16xi1>, vector<16xf32>
}

// -----

func.func @scan_reduction_dim_constraint(%arg0: vector<2x3xi32>, %arg1: vector<3xi32>) -> vector<3xi32> {
// expected-error@+1 {{'vector.scan' op reduction dimension 5 has to be less than 2}}
%0:2 = vector.scan <add>, %arg0, %arg1 {inclusive = true, reduction_dim = 5} :
Expand Down Expand Up @@ -1952,15 +2016,15 @@ func.func @vector_load(%src : memref<?xi8>) {

// -----

func.func @invalid_load_alignment(%memref: memref<4xi32>, %c0: index) {
func.func @load_negative_alignment(%memref: memref<4xi32>, %c0: index) {
// expected-error @below {{'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}}
%val = vector.load %memref[%c0] { alignment = -1 } : memref<4xi32>, vector<4xi32>
return
}

// -----

func.func @invalid_load_alignment(%memref: memref<4xi32>, %c0: index) {
func.func @load_non_pow_of_2_alignment(%memref: memref<4xi32>, %c0: index) {
// expected-error @below {{'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}}
%val = vector.load %memref[%c0] { alignment = 3 } : memref<4xi32>, vector<4xi32>
return
Expand All @@ -1981,15 +2045,15 @@ func.func @vector_store(%dest : memref<?xi8>, %vec : vector<16x16xi8>) {

// -----

func.func @invalid_store_alignment(%memref: memref<4xi32>, %val: vector<4xi32>, %c0: index) {
func.func @store_negative_alignment(%memref: memref<4xi32>, %val: vector<4xi32>, %c0: index) {
// expected-error @below {{'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}}
vector.store %val, %memref[%c0] { alignment = -1 } : memref<4xi32>, vector<4xi32>
return
}

// -----

func.func @invalid_store_alignment(%memref: memref<4xi32>, %val: vector<4xi32>, %c0: index) {
func.func @store_non_pow_of_2_alignment(%memref: memref<4xi32>, %val: vector<4xi32>, %c0: index) {
// expected-error @below {{'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}}
vector.store %val, %memref[%c0] { alignment = 3 } : memref<4xi32>, vector<4xi32>
return
Expand Down