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

Open
wants to merge 11 commits into
base: main
Choose a base branch
from
Open

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

Show file tree
Hide file tree
Changes from 6 commits
Commits
Show all changes
11 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
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
112 changes: 104 additions & 8 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 @@ -1382,6 +1382,11 @@ def Vector_TransferReadOp :
An additional `1` broadcast is required. On a GPU this broadcast could be
implemented using a warp-shuffle if loop `j` were mapped to `threadIdx.x`.

An optional `alignment` attribute allows to specify the byte alignment of the
Copy link
Member

Choose a reason for hiding this comment

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

This seems to be added under the wrong op

Copy link
Contributor Author

Choose a reason for hiding this comment

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

Thanks @kuhar, I removed it here 3180cd0 and added it back to vector.gather here: 47db5b1

gather 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.

Syntax
```
operation ::= ssa-id `=` `vector.transfer_read` ssa-use-list
Expand Down Expand Up @@ -1714,7 +1719,6 @@ def Vector_LoadOp : Vector_Op<"load", [
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.
}];

let arguments = (ins Arg<AnyMemRef, "the reference to load from",
Expand Down Expand Up @@ -1830,7 +1834,6 @@ def Vector_StoreOp : Vector_Op<"store", [
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.
}];

let arguments = (ins
Expand Down Expand Up @@ -1919,7 +1922,6 @@ def Vector_MaskedLoadOp :
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.
Copy link
Contributor

Choose a reason for hiding this comment

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

Can you remind me what happens when alignment is not specified?

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 originally wanted thought about removing this line since I imagined that the constructors using the llvm::Maybe align will be preferred, but I now believe that adding this line back makes more sense since there are other constructors as well and the actual value stored is an integer attribute. Thanks for pointing it out! 47db5b1

Copy link
Contributor

Choose a reason for hiding this comment

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

And just to double check - is 0 the default value?

Copy link
Contributor Author

@amd-eochoalo amd-eochoalo Aug 12, 2025 •
edited
Loading

Choose a reason for hiding this comment

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

0 is not the default value of the alignment attribute in the operation. I think there's some ambiguity here regarding what "default" means. It could mean two things in my opinion:

  1. The default parameter to one of the constructors.
  2. The default value in the Operation's alignment field.

In the original PR (#144344) the default parameter to one of the constructors is indeed zero, but the attribute is optional and the attribute linked to the operation is actually a nullptr.

                   CArg<"uint64_t", "0">:$alignment), [{
      return build($_builder, $_state, memref, indices, nontemporal,
                   alignment != 0 ? $_builder.getI64IntegerAttr(alignment) :
                                    nullptr);

In PR #151690 the default parameter for these constructors was changed from a uint64_t type to an llvm::MaybeAlign, but keeps the attribute linked to the operation to be nullptr when there is no alignment requirement. I.e., PR #15169

  1. The default parameter to these constructors is llvm::MaybeAlign()
  2. When the default parameter is llvm::MaybeAlign() the integer attribute pointer is nullptr.

Just to be complete, I believe in both cases the operation could have the field could be nullptr or point to I64IntegerAttr(0) to indicate no alignment requirement. For example, if the user used a different constructor passing all attributes in order.

I think having the documentation indicate that a value of zero indicates no specific alignment requirements is still correct as the Operation's alignment field is still an integer (when present) and it being zero would still signifies no specific alignment requirements.

I think we could also make the alignment attribute required by removing the OptionalAttr and then setting the alignment field point to I64IntegerAttr(0) to remove the nullptr and have solely I64IntegerAttr(0) mean no specific alignment requirement. Happy to add changes if you think it is required :-).

I could also change the line to say that a value of llvm::MaybeAlign() indicates no specific alignment requirements and values of llvm::Align(n) for n bigger than zero to be alignment requirements.

}];
let extraClassDeclaration = [{
MemRefType getMemRefType() {
Expand Down Expand Up @@ -2012,7 +2014,6 @@ def Vector_MaskedStoreOp :
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.
}];
let extraClassDeclaration = [{
MemRefType getMemRefType() {
Expand Down Expand Up @@ -2054,7 +2055,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 @@ -2111,6 +2114,31 @@ 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);
}]>,
OpBuilder<(ins "TypeRange":$resultTypes,
"Value":$base,
"ValueRange":$indices,
"Value":$index_vec,
"Value":$mask,
"Value":$passthrough,
CArg<"llvm::MaybeAlign", "llvm::MaybeAlign()">:$alignment), [{
return build($_builder, $_state, resultTypes, base, indices, index_vec, mask, passthrough,
alignment.has_value() ? $_builder.getI64IntegerAttr(alignment->value()) :
nullptr);
}]>
];
}

def Vector_ScatterOp :
Expand All @@ -2119,7 +2147,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 +2183,11 @@ 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.

Examples:

```mlir
Expand All @@ -2177,14 +2212,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 +2266,11 @@ 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.

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

Examples:
Expand Down Expand Up @@ -2246,14 +2301,39 @@ 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);
}]>,
OpBuilder<(ins "TypeRange":$resultTypes,
"Value":$base,
"ValueRange":$indices,
"Value":$mask,
"Value":$passthrough,
CArg<"llvm::MaybeAlign", "llvm::MaybeAlign()">:$alignment), [{
return build($_builder, $_state, resultTypes, 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 +2364,11 @@ 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.

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

Examples:
Expand Down Expand Up @@ -2312,6 +2397,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
64 changes: 64 additions & 0 deletions mlir/test/Dialect/Vector/invalid.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1470,6 +1470,24 @@ func.func @gather_pass_thru_type_mismatch(%base: memref<?xf32>, %indices: vector

// -----

func.func @gather_invalid_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_invalid_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_invalid_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_invalid_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_invalid_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_invalid_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_invalid_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_invalid_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