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

amd-eochoalo · 2025-08-07T14:12:03Z

Following #144344, #152207, #151690, this PR adds the alignment attribute to the following operations in the vector dialect:

compressstore
expandload
vector.scatter
vector.gather

llvmbot · 2025-08-08T17:56:23Z

@llvm/pr-subscribers-mlir-vector

@llvm/pr-subscribers-mlir

Author: Erick Ochoa Lopez (amd-eochoalo)

Changes

Following #144344, #152207, #151690, this PR adds the alignment attribute to the following operations in the vector dialect:

compressstore
expandload
vector.scatter
vector.gather

Full diff: https://github.com/llvm/llvm-project/pull/152507.diff

2 Files Affected:

(modified) mlir/include/mlir/Dialect/Vector/IR/VectorOps.td (+84-4)
(modified) mlir/test/Dialect/Vector/invalid.mlir (+64)

diff --git a/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td b/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
index b3b8afdd8b4c1..aae2051600251 100644
--- a/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
+++ b/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
@@ -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 = [{
@@ -2111,6 +2113,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 :
@@ -2119,7 +2146,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
@@ -2177,6 +2206,19 @@ 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 :
@@ -2184,7 +2226,9 @@ def Vector_ExpandLoadOp :
     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";
@@ -2246,6 +2290,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);
+    }]>,
+    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 :
@@ -2253,7 +2320,9 @@ def Vector_CompressStoreOp :
     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";
 
@@ -2312,6 +2381,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 :
diff --git a/mlir/test/Dialect/Vector/invalid.mlir b/mlir/test/Dialect/Vector/invalid.mlir
index 211e16db85a94..68b07ec82aeb7 100644
--- a/mlir/test/Dialect/Vector/invalid.mlir
+++ b/mlir/test/Dialect/Vector/invalid.mlir
@@ -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
@@ -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}}
@@ -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}}
@@ -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} :

kuhar

Can you update op documentation and describe the semantics?

banach-space

Thanks!

mlir/test/Dialect/Vector/invalid.mlir

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

banach-space · 2025-08-11T09:12:16Z

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

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


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

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

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

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:

The default parameter to one of the constructors.

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

The default parameter to these constructors is llvm::MaybeAlign()

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.

kuhar · 2025-08-11T13:30:04Z

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

@@ -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


This seems to be added under the wrong op

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

This is a squash of PR llvm#152507

amd-eochoalo force-pushed the eochoa/2025-08-06/vector-ops-alignment branch 2 times, most recently from 5a47fff to e2ad0f9 Compare August 8, 2025 17:42

amd-eochoalo marked this pull request as ready for review August 8, 2025 17:55

amd-eochoalo requested a review from kuhar as a code owner August 8, 2025 17:55

llvmbot added mlir:vectorops mlir mlir:vector labels Aug 8, 2025

kuhar reviewed Aug 8, 2025

View reviewed changes

amd-eochoalo added 6 commits August 8, 2025 18:19

[mlir][vector] Add alignment to vector.gather.

a97cbc6

[mlir][vector] Add alignment to vector.scatter

158b91c

[mlir][vector] Add alignment to compressstore

36949d1

[mlir][vector] Add alignment to expandload

a4d820f

Use llvm::MaybeAlign

e2ad0f9

Add documentation for alignment attribute

92b3886

banach-space reviewed Aug 11, 2025

View reviewed changes

kuhar reviewed Aug 11, 2025

View reviewed changes

amd-eochoalo added 5 commits August 11, 2025 14:07

Change names in test functions

482ad75

Remove unnecessary constructor

9b451db

Remove documentation from wrong op

3180cd0

Fix documention

47db5b1

Remove unnecessary constructor

a714521

amd-eochoalo added a commit to amd-eochoalo/llvm-project that referenced this pull request Aug 13, 2025

[mlir][vector] Add alignment to vector operations.

a0cc39a

This is a squash of PR llvm#152507

amd-eochoalo mentioned this pull request Aug 13, 2025

[mlir][vector] Propagate alignment from vector to llvm dialects. #153482

Draft

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

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

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

amd-eochoalo commented Aug 7, 2025 •

edited

Loading

Uh oh!

llvmbot commented Aug 8, 2025 •

edited

Loading

Uh oh!

kuhar left a comment

Uh oh!

banach-space left a comment

Uh oh!

Uh oh!

Uh oh!

banach-space Aug 11, 2025

Uh oh!

amd-eochoalo Aug 11, 2025

Uh oh!

banach-space Aug 11, 2025

Uh oh!

amd-eochoalo Aug 12, 2025 •

edited

Loading

Uh oh!

kuhar Aug 11, 2025

Uh oh!

amd-eochoalo Aug 11, 2025

Uh oh!

Uh oh!

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

Are you sure you want to change the base?

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

Conversation

amd-eochoalo commented Aug 7, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

llvmbot commented Aug 8, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

kuhar left a comment

Choose a reason for hiding this comment

Uh oh!

banach-space left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

banach-space Aug 11, 2025

Choose a reason for hiding this comment

Uh oh!

amd-eochoalo Aug 11, 2025

Choose a reason for hiding this comment

Uh oh!

banach-space Aug 11, 2025

Choose a reason for hiding this comment

Uh oh!

amd-eochoalo Aug 12, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

kuhar Aug 11, 2025

Choose a reason for hiding this comment

Uh oh!

amd-eochoalo Aug 11, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

amd-eochoalo commented Aug 7, 2025 •

edited

Loading

llvmbot commented Aug 8, 2025 •

edited

Loading

amd-eochoalo Aug 12, 2025 •

edited

Loading