[Mlir-commits] [mlir] [mlir][vector] Add alignment attribute to vector operations. (PR #152507)
Erick Ochoa Lopez
llvmlistbot at llvm.org
Thu Aug 7 07:12:14 PDT 2025
https://github.com/amd-eochoalo created https://github.com/llvm/llvm-project/pull/152507
Following #144344, #152207, #151690, this PR adds the alignment attribute to the following operations in the vector dialect:
* `compressstore`
* `expandload`
* `transfer_read`
* `transfer_write`
* `vector.scatter`
* `vector.gather`
>From 394c64fa315069c9db524ddc358b8dbe5aae72cf Mon Sep 17 00:00:00 2001
From: Erick Ochoa <erick.ochoalopez at amd.com>
Date: Wed, 6 Aug 2025 16:34:02 -0700
Subject: [PATCH 1/7] Use llvm::Align
---
mlir/include/mlir/Dialect/Vector/IR/VectorOps.h | 1 +
mlir/include/mlir/Dialect/Vector/IR/VectorOps.td | 12 ++++++------
2 files changed, 7 insertions(+), 6 deletions(-)
diff --git a/mlir/include/mlir/Dialect/Vector/IR/VectorOps.h b/mlir/include/mlir/Dialect/Vector/IR/VectorOps.h
index 364c1728715e8..63410b8bea747 100644
--- a/mlir/include/mlir/Dialect/Vector/IR/VectorOps.h
+++ b/mlir/include/mlir/Dialect/Vector/IR/VectorOps.h
@@ -32,6 +32,7 @@
#include "mlir/Interfaces/ViewLikeInterface.h"
#include "llvm/ADT/SetVector.h"
#include "llvm/ADT/StringExtras.h"
+#include "llvm/Support/Alignment.h"
// Pull in all enum type definitions and utility function declarations.
#include "mlir/Dialect/Vector/IR/VectorEnums.h.inc"
diff --git a/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td b/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
index dc55704c36183..eeedf68a1df7c 100644
--- a/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
+++ b/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
@@ -1729,18 +1729,18 @@ def Vector_LoadOp : Vector_Op<"load", [
"Value":$base,
"ValueRange":$indices,
CArg<"bool", "false">:$nontemporal,
- CArg<"uint64_t", "0">:$alignment), [{
+ CArg<"llvm::Align", "llvm::Align()">:$alignment), [{
return build($_builder, $_state, resultType, base, indices, nontemporal,
- alignment != 0 ? $_builder.getI64IntegerAttr(alignment) :
+ alignment != llvm::Align() ? $_builder.getI64IntegerAttr(alignment.value()) :
nullptr);
}]>,
OpBuilder<(ins "TypeRange":$resultTypes,
"Value":$base,
"ValueRange":$indices,
CArg<"bool", "false">:$nontemporal,
- CArg<"uint64_t", "0">:$alignment), [{
+ CArg<"llvm::Align", "llvm::Align()">:$alignment), [{
return build($_builder, $_state, resultTypes, base, indices, nontemporal,
- alignment != 0 ? $_builder.getI64IntegerAttr(alignment) :
+ alignment != llvm::Align() ? $_builder.getI64IntegerAttr(alignment.value()) :
nullptr);
}]>
];
@@ -1847,9 +1847,9 @@ def Vector_StoreOp : Vector_Op<"store", [
"Value":$base,
"ValueRange":$indices,
CArg<"bool", "false">:$nontemporal,
- CArg<"uint64_t", "0">:$alignment), [{
+ CArg<"llvm::Align", "llvm::Align()">:$alignment), [{
return build($_builder, $_state, valueToStore, base, indices, nontemporal,
- alignment != 0 ? $_builder.getI64IntegerAttr(alignment) :
+ alignment != llvm::Align() ? $_builder.getI64IntegerAttr(alignment.value()) :
nullptr);
}]>
];
>From da405541dd57a50b38a9cf5a862b1cddc1bc7d19 Mon Sep 17 00:00:00 2001
From: Erick Ochoa <erick.ochoalopez at amd.com>
Date: Wed, 6 Aug 2025 19:48:16 -0700
Subject: [PATCH 2/7] [mlir][vector] Add alignment to vector.gather.
---
.../mlir/Dialect/Vector/IR/VectorOps.td | 29 ++++++++++++++++++-
mlir/test/Dialect/Vector/invalid.mlir | 18 ++++++++++++
2 files changed, 46 insertions(+), 1 deletion(-)
diff --git a/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td b/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
index eeedf68a1df7c..f2b55be19e192 100644
--- a/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
+++ b/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
@@ -2003,7 +2003,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 = [{
@@ -2060,6 +2062,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::Align", "llvm::Align()">:$alignment), [{
+ return build($_builder, $_state, resultType, base, indices, index_vec, mask, passthrough,
+ alignment != llvm::Align() ? $_builder.getI64IntegerAttr(alignment.value()) :
+ nullptr);
+ }]>,
+ OpBuilder<(ins "TypeRange":$resultTypes,
+ "Value":$base,
+ "ValueRange":$indices,
+ "Value":$index_vec,
+ "Value":$mask,
+ "Value":$passthrough,
+ CArg<"llvm::Align", "llvm::Align()">:$alignment), [{
+ return build($_builder, $_state, resultTypes, base, indices, index_vec, mask, passthrough,
+ alignment != llvm::Align() ? $_builder.getI64IntegerAttr(alignment.value()) :
+ nullptr);
+ }]>
+ ];
}
def Vector_ScatterOp :
diff --git a/mlir/test/Dialect/Vector/invalid.mlir b/mlir/test/Dialect/Vector/invalid.mlir
index c21de562d05e1..6a4c7a5623a43 100644
--- a/mlir/test/Dialect/Vector/invalid.mlir
+++ b/mlir/test/Dialect/Vector/invalid.mlir
@@ -1430,6 +1430,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 at +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 at +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
>From 12eead5b3dd6496afe11653be4d03e75c39f7356 Mon Sep 17 00:00:00 2001
From: Erick Ochoa <erick.ochoalopez at amd.com>
Date: Wed, 6 Aug 2025 19:58:43 -0700
Subject: [PATCH 3/7] [mlir][vector] Add alignment to vector.scatter
---
.../mlir/Dialect/Vector/IR/VectorOps.td | 17 ++++++++++++++++-
mlir/test/Dialect/Vector/invalid.mlir | 18 ++++++++++++++++++
2 files changed, 34 insertions(+), 1 deletion(-)
diff --git a/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td b/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
index f2b55be19e192..1ad527d5619d3 100644
--- a/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
+++ b/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
@@ -2095,7 +2095,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
@@ -2153,6 +2155,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::Align", "llvm::Align()">: $alignment), [{
+ return build($_builder, $_state, base, indices, index_vec, mask, valueToStore,
+ alignment != llvm::Align() ? $_builder.getI64IntegerAttr(alignment.value()) :
+ nullptr);
+ }]>
+ ];
}
def Vector_ExpandLoadOp :
diff --git a/mlir/test/Dialect/Vector/invalid.mlir b/mlir/test/Dialect/Vector/invalid.mlir
index 6a4c7a5623a43..c5f5c2694b6ff 100644
--- a/mlir/test/Dialect/Vector/invalid.mlir
+++ b/mlir/test/Dialect/Vector/invalid.mlir
@@ -1509,6 +1509,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 at +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 at +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 at +1 {{'vector.expandload' op base and result element type should match}}
>From 2b512366d25a7e417dd4b5f5e1ea46b320d481ae Mon Sep 17 00:00:00 2001
From: Erick Ochoa <erick.ochoalopez at amd.com>
Date: Wed, 6 Aug 2025 20:48:21 -0700
Subject: [PATCH 4/7] [mlir][vector] Add alignment to compressstore
---
mlir/include/mlir/Dialect/Vector/IR/VectorOps.td | 15 ++++++++++++++-
mlir/test/Dialect/Vector/invalid.mlir | 14 ++++++++++++++
2 files changed, 28 insertions(+), 1 deletion(-)
diff --git a/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td b/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
index 1ad527d5619d3..d3c3260ae4484 100644
--- a/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
+++ b/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
@@ -2244,7 +2244,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";
@@ -2303,6 +2305,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::Align", "llvm::Align()">:$alignment), [{
+ return build($_builder, $_state, base, indices, valueToStore, mask,
+ alignment != llvm::Align() ? $_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 c5f5c2694b6ff..23c908a95616f 100644
--- a/mlir/test/Dialect/Vector/invalid.mlir
+++ b/mlir/test/Dialect/Vector/invalid.mlir
@@ -1599,6 +1599,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 at +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} :
>From 5a47ffff57cf05ed0b86ca232f3aa23b9c168111 Mon Sep 17 00:00:00 2001
From: Erick Ochoa <erick.ochoalopez at amd.com>
Date: Wed, 6 Aug 2025 20:57:52 -0700
Subject: [PATCH 5/7] [mlir][vector] Add alignment to expandload
---
.../mlir/Dialect/Vector/IR/VectorOps.td | 27 ++++++++++++++++++-
mlir/test/Dialect/Vector/invalid.mlir | 14 ++++++++++
2 files changed, 40 insertions(+), 1 deletion(-)
diff --git a/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td b/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
index d3c3260ae4484..62fac6ef2ea36 100644
--- a/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
+++ b/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
@@ -2175,7 +2175,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";
@@ -2237,6 +2239,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::Align", "llvm::Align()">:$alignment), [{
+ return build($_builder, $_state, resultType, base, indices, mask, passthrough,
+ alignment != llvm::Align() ? $_builder.getI64IntegerAttr(alignment.value()) :
+ nullptr);
+ }]>,
+ OpBuilder<(ins "TypeRange":$resultTypes,
+ "Value":$base,
+ "ValueRange":$indices,
+ "Value":$mask,
+ "Value":$passthrough,
+ CArg<"llvm::Align", "llvm::Align()">:$alignment), [{
+ return build($_builder, $_state, resultTypes, base, indices, mask, passthrough,
+ alignment != llvm::Align() ? $_builder.getI64IntegerAttr(alignment.value()) :
+ nullptr);
+ }]>
+ ];
}
def Vector_CompressStoreOp :
diff --git a/mlir/test/Dialect/Vector/invalid.mlir b/mlir/test/Dialect/Vector/invalid.mlir
index 23c908a95616f..b95df0fe30eb2 100644
--- a/mlir/test/Dialect/Vector/invalid.mlir
+++ b/mlir/test/Dialect/Vector/invalid.mlir
@@ -1567,6 +1567,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 at +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 at +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 at +1 {{'vector.compressstore' op base and valueToStore element type should match}}
>From 9fdfd0371a723cd69a0ab7fe58c0ceccfa8d02cd Mon Sep 17 00:00:00 2001
From: Erick Ochoa <erick.ochoalopez at amd.com>
Date: Wed, 6 Aug 2025 23:10:21 -0700
Subject: [PATCH 6/7] [mlir][vector] Add alignment to transfer_read
---
.../mlir/Dialect/Vector/IR/VectorOps.td | 38 +++++++++++++++++--
mlir/lib/Dialect/Vector/IR/VectorOps.cpp | 16 +++++---
mlir/test/Dialect/Vector/invalid.mlir | 16 ++++++++
3 files changed, 60 insertions(+), 10 deletions(-)
diff --git a/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td b/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
index 62fac6ef2ea36..bef40556c4905 100644
--- a/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
+++ b/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
@@ -1253,7 +1253,9 @@ def Vector_TransferReadOp :
AffineMapAttr:$permutation_map,
AnyType:$padding,
Optional<VectorOfNonZeroRankOf<[I1]>>:$mask,
- BoolArrayAttr:$in_bounds)>,
+ BoolArrayAttr:$in_bounds,
+ ConfinedAttr<OptionalAttr<I64Attr>,
+ [AllAttrOf<[IntPositive, IntPowerOf2]>]>:$alignment)>,
Results<(outs AnyVectorOfAnyRank:$vector)> {
let summary = "Reads a supervector from memory into an SSA vector value.";
@@ -1443,6 +1445,31 @@ def Vector_TransferReadOp :
}];
let builders = [
+ OpBuilder<(ins "Type":$resultType,
+ "Value":$base,
+ "ValueRange":$indices,
+ "AffineMapAttr":$permutationMapAttr,
+ "Value":$padding,
+ "Value":$mask,
+ "ArrayAttr":$inBoundsAttr,
+ CArg<"llvm::Align", "llvm::Align()">:$alignment), [{
+ return build($_builder, $_state, resultType, base, indices, permutationMapAttr, padding, mask, inBoundsAttr,
+ alignment != llvm::Align() ? $_builder.getI64IntegerAttr(alignment.value()) :
+ nullptr);
+ }]>,
+ OpBuilder<(ins "Type":$resultType,
+ "Value":$base,
+ "ValueRange":$indices,
+ "AffineMap":$permutationMap,
+ "Value":$padding,
+ "Value":$mask,
+ "ArrayAttr":$inBoundsAttr,
+ CArg<"llvm::Align", "llvm::Align()">:$alignment), [{
+ auto permutationMapAttr = AffineMapAttr::get(permutationMap);
+ return build($_builder, $_state, resultType, base, indices, permutationMapAttr, padding, mask, inBoundsAttr,
+ alignment != llvm::Align() ? $_builder.getI64IntegerAttr(alignment.value()) :
+ nullptr);
+ }]>,
/// 1. Builder that sets padding to `padding` or poison if not provided and
/// an empty mask (variant with attrs).
OpBuilder<(ins "VectorType":$vectorType,
@@ -1450,7 +1477,8 @@ def Vector_TransferReadOp :
"ValueRange":$indices,
"std::optional<Value>":$padding,
"AffineMapAttr":$permutationMapAttr,
- "ArrayAttr":$inBoundsAttr)>,
+ "ArrayAttr":$inBoundsAttr,
+ CArg<"llvm::Align", "llvm::Align()">:$alignment)>,
/// 2. Builder that sets padding to `padding` or poison if not provided and
/// an empty mask (variant without attrs).
OpBuilder<(ins "VectorType":$vectorType,
@@ -1458,14 +1486,16 @@ def Vector_TransferReadOp :
"ValueRange":$indices,
"std::optional<Value>":$padding,
"AffineMap":$permutationMap,
- CArg<"std::optional<ArrayRef<bool>>", "::std::nullopt">:$inBounds)>,
+ CArg<"std::optional<ArrayRef<bool>>", "::std::nullopt">:$inBounds,
+ CArg<"llvm::Align", "llvm::Align()">:$alignment)>,
/// 3. Builder that sets padding to `padding` or poison if not provided and
/// permutation map to 'getMinorIdentityMap'.
OpBuilder<(ins "VectorType":$vectorType,
"Value":$source,
"ValueRange":$indices,
"std::optional<Value>":$padding,
- CArg<"std::optional<ArrayRef<bool>>", "::std::nullopt">:$inBounds)>
+ CArg<"std::optional<ArrayRef<bool>>", "::std::nullopt">:$inBounds,
+ CArg<"llvm::Align", "llvm::Align()">:$alignment)>,
];
let extraClassDeclaration = [{
diff --git a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
index a450056a3041a..9bff8261a0d5d 100644
--- a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
+++ b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
@@ -43,6 +43,7 @@
#include "llvm/ADT/SmallVector.h"
#include "llvm/ADT/StringSet.h"
#include "llvm/ADT/TypeSwitch.h"
+#include "llvm/Support/Alignment.h"
#include "llvm/Support/Casting.h"
#include <cassert>
@@ -4320,13 +4321,14 @@ void TransferReadOp::build(OpBuilder &builder, OperationState &result,
VectorType vectorType, Value source,
ValueRange indices, std::optional<Value> padding,
AffineMapAttr permutationMapAttr,
- /*optional*/ ArrayAttr inBoundsAttr) {
+ /*optional*/ ArrayAttr inBoundsAttr,
+ /*optional*/ llvm::Align alignment) {
Type elemType = llvm::cast<ShapedType>(source.getType()).getElementType();
if (!padding)
padding = ub::PoisonOp::create(builder, result.location, elemType);
build(builder, result, vectorType, source, indices, permutationMapAttr,
- *padding, /*mask=*/Value(), inBoundsAttr);
+ *padding, /*mask=*/Value(), inBoundsAttr, alignment);
}
/// 2. Builder that sets padding to zero an empty mask (variant without attrs).
@@ -4334,7 +4336,8 @@ void TransferReadOp::build(OpBuilder &builder, OperationState &result,
VectorType vectorType, Value source,
ValueRange indices, std::optional<Value> padding,
AffineMap permutationMap,
- std::optional<ArrayRef<bool>> inBounds) {
+ std::optional<ArrayRef<bool>> inBounds,
+ /*optional*/ llvm::Align alignment) {
auto permutationMapAttr = AffineMapAttr::get(permutationMap);
auto inBoundsAttr = (inBounds && !inBounds.value().empty())
? builder.getBoolArrayAttr(inBounds.value())
@@ -4344,14 +4347,15 @@ void TransferReadOp::build(OpBuilder &builder, OperationState &result,
if (!padding)
padding = ub::PoisonOp::create(builder, result.location, elemType);
build(builder, result, vectorType, source, indices, *padding,
- permutationMapAttr, inBoundsAttr);
+ permutationMapAttr, inBoundsAttr, alignment);
}
/// 3. Builder that sets permutation map to 'getMinorIdentityMap'.
void TransferReadOp::build(OpBuilder &builder, OperationState &result,
VectorType vectorType, Value source,
ValueRange indices, std::optional<Value> padding,
- std::optional<ArrayRef<bool>> inBounds) {
+ std::optional<ArrayRef<bool>> inBounds,
+ /*optional*/ llvm::Align alignment) {
AffineMap permutationMap = getTransferMinorIdentityMap(
llvm::cast<ShapedType>(source.getType()), vectorType);
auto permutationMapAttr = AffineMapAttr::get(permutationMap);
@@ -4364,7 +4368,7 @@ void TransferReadOp::build(OpBuilder &builder, OperationState &result,
padding = ub::PoisonOp::create(builder, result.location, elemType);
build(builder, result, vectorType, source, indices, permutationMapAttr,
*padding,
- /*mask=*/Value(), inBoundsAttr);
+ /*mask=*/Value(), inBoundsAttr, alignment);
}
template <typename EmitFun>
diff --git a/mlir/test/Dialect/Vector/invalid.mlir b/mlir/test/Dialect/Vector/invalid.mlir
index b95df0fe30eb2..16295fc23632c 100644
--- a/mlir/test/Dialect/Vector/invalid.mlir
+++ b/mlir/test/Dialect/Vector/invalid.mlir
@@ -471,6 +471,22 @@ func.func @test_vector.transfer_read(%arg0: memref<?xindex>) -> vector<3x4xindex
// -----
+func.func @test_vector.transfer_read(%arg0: tensor<f32>, %f0: f32) {
+ // expected-error at +1 {{'vector.transfer_read' 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.transfer_read %arg0[], %f0 {permutation_map = affine_map<()->(0)>, alignment = -1} : tensor<f32>, vector<1xf32>
+ return
+}
+
+// -----
+
+func.func @test_vector.transfer_read(%arg0: tensor<f32>, %f0: f32) {
+ // expected-error at +1 {{'vector.transfer_read' 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.transfer_read %arg0[], %f0 {permutation_map = affine_map<()->(0)>, alignment = 3} : tensor<f32>, vector<1xf32>
+ return
+}
+
+// -----
+
func.func @test_vector.transfer_write(%arg0: memref<?xvector<2xindex>>) {
%c3 = arith.constant 3 : index
// expected-error at +1 {{expected a custom permutation_map when rank(source) != rank(destination)}}
>From 4b3593b68e7afffbddc75c311aade0e6cfd328e5 Mon Sep 17 00:00:00 2001
From: Erick Ochoa <erick.ochoalopez at amd.com>
Date: Thu, 7 Aug 2025 14:59:25 -0700
Subject: [PATCH 7/7] [mlir][vector] Add alignment for transfer_write
---
.../mlir/Dialect/Vector/IR/VectorOps.td | 41 ++++++++++++++++---
mlir/lib/Dialect/Vector/IR/VectorOps.cpp | 24 +++++++----
mlir/test/Dialect/Vector/invalid.mlir | 16 ++++++++
3 files changed, 68 insertions(+), 13 deletions(-)
diff --git a/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td b/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
index bef40556c4905..5777e8fbe2e8c 100644
--- a/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
+++ b/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
@@ -1530,7 +1530,9 @@ def Vector_TransferWriteOp :
Variadic<Index>:$indices,
AffineMapAttr:$permutation_map,
Optional<VectorOfNonZeroRankOf<[I1]>>:$mask,
- BoolArrayAttr:$in_bounds)>,
+ BoolArrayAttr:$in_bounds,
+ ConfinedAttr<OptionalAttr<I64Attr>,
+ [AllAttrOf<[IntPositive, IntPowerOf2]>]>:$alignment)>,
Results<(outs Optional<AnyRankedTensor>:$result)> {
let summary = "The vector.transfer_write op writes a supervector to memory.";
@@ -1635,31 +1637,60 @@ def Vector_TransferWriteOp :
}];
let builders = [
+ OpBuilder<(ins "Type":$resultType,
+ "Value":$vector,
+ "Value":$dest,
+ "ValueRange":$indices,
+ "AffineMapAttr":$permutationMapAttr,
+ "Value":$mask,
+ "ArrayAttr":$inBoundsAttr,
+ CArg<"llvm::Align", "llvm::Align()">:$alignment), [{
+ return build($_builder, $_state, resultType, vector, dest, indices, permutationMapAttr, mask, inBoundsAttr,
+ alignment != llvm::Align() ? $_builder.getI64IntegerAttr(alignment.value()) :
+ nullptr);
+ }]>,
+ OpBuilder<(ins "Type":$resultType,
+ "Value":$vector,
+ "Value":$dest,
+ "ValueRange":$indices,
+ "AffineMap":$permutationMap,
+ "Value":$mask,
+ "ArrayAttr":$inBoundsAttr,
+ CArg<"llvm::Align", "llvm::Align()">:$alignment), [{
+ auto permutationMapAttr = AffineMapAttr::get(permutationMap);
+ return build($_builder, $_state, resultType, vector, dest, indices, permutationMapAttr, mask, inBoundsAttr,
+ alignment != llvm::Align() ? $_builder.getI64IntegerAttr(alignment.value()) :
+ nullptr);
+ }]>,
/// 1. Builder with type inference.
OpBuilder<(ins "Value":$vector,
"Value":$dest,
"ValueRange":$indices,
"AffineMapAttr":$permutationMapAttr,
"Value":$mask,
- "ArrayAttr":$inBoundsAttr)>,
+ "ArrayAttr":$inBoundsAttr,
+ CArg<"llvm::Align", "llvm::Align()">:$alignment)>,
/// 2. Builder with type inference that sets an empty mask (variant with attrs).
OpBuilder<(ins "Value":$vector,
"Value":$dest,
"ValueRange":$indices,
"AffineMapAttr":$permutationMapAttr,
- "ArrayAttr":$inBoundsAttr)>,
+ "ArrayAttr":$inBoundsAttr,
+ CArg<"llvm::Align", "llvm::Align()">:$alignment)>,
/// 3. Builder with type inference that sets an empty mask (variant without attrs).
OpBuilder<(ins "Value":$vector,
"Value":$dest,
"ValueRange":$indices,
"AffineMap":$permutationMap,
- CArg<"std::optional<ArrayRef<bool>>", "::std::nullopt">:$inBounds)>,
+ CArg<"std::optional<ArrayRef<bool>>", "::std::nullopt">:$inBounds,
+ CArg<"llvm::Align", "llvm::Align()">:$alignment)>,
/// 4. Builder with type inference that sets an empty mask and sets permutation
/// map to 'getMinorIdentityMap'.
OpBuilder<(ins "Value":$vector,
"Value":$dest,
"ValueRange":$indices,
- CArg<"std::optional<ArrayRef<bool>>", "::std::nullopt">:$inBounds)>,
+ CArg<"std::optional<ArrayRef<bool>>", "::std::nullopt">:$inBounds,
+ CArg<"llvm::Align", "llvm::Align()">:$alignment)>,
];
let extraClassDeclaration = [{
diff --git a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
index 9bff8261a0d5d..cd56f888b3ac1 100644
--- a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
+++ b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
@@ -4913,19 +4913,24 @@ void TransferWriteOp::build(OpBuilder &builder, OperationState &result,
Value vector, Value dest, ValueRange indices,
AffineMapAttr permutationMapAttr,
/*optional*/ Value mask,
- /*optional*/ ArrayAttr inBoundsAttr) {
+ /*optional*/ ArrayAttr inBoundsAttr,
+ /*optional*/ llvm::Align alignment) {
Type resultType = llvm::dyn_cast<RankedTensorType>(dest.getType());
build(builder, result, resultType, vector, dest, indices, permutationMapAttr,
- mask, inBoundsAttr);
+ mask, inBoundsAttr,
+ alignment != llvm::Align()
+ ? builder.getI64IntegerAttr(alignment.value())
+ : nullptr);
}
/// 2. Builder with type inference that sets an empty mask (variant with attrs).
void TransferWriteOp::build(OpBuilder &builder, OperationState &result,
Value vector, Value dest, ValueRange indices,
AffineMapAttr permutationMapAttr,
- /*optional*/ ArrayAttr inBoundsAttr) {
+ /*optional*/ ArrayAttr inBoundsAttr,
+ /*optional*/ llvm::Align alignment) {
build(builder, result, vector, dest, indices, permutationMapAttr,
- /*mask=*/Value(), inBoundsAttr);
+ /*mask=*/Value(), inBoundsAttr, alignment);
}
/// 3. Builder with type inference that sets an empty mask (variant without
@@ -4933,7 +4938,8 @@ void TransferWriteOp::build(OpBuilder &builder, OperationState &result,
void TransferWriteOp::build(OpBuilder &builder, OperationState &result,
Value vector, Value dest, ValueRange indices,
AffineMap permutationMap,
- std::optional<ArrayRef<bool>> inBounds) {
+ std::optional<ArrayRef<bool>> inBounds,
+ /*optional*/ llvm::Align alignment) {
auto permutationMapAttr = AffineMapAttr::get(permutationMap);
auto inBoundsAttr =
(inBounds && !inBounds.value().empty())
@@ -4941,18 +4947,20 @@ void TransferWriteOp::build(OpBuilder &builder, OperationState &result,
: builder.getBoolArrayAttr(SmallVector<bool>(
llvm::cast<VectorType>(vector.getType()).getRank(), false));
build(builder, result, vector, dest, indices, permutationMapAttr,
- /*mask=*/Value(), inBoundsAttr);
+ /*mask=*/Value(), inBoundsAttr, alignment);
}
/// 4. Builder with type inference that sets an empty mask and sets permutation
/// map to 'getMinorIdentityMap'.
void TransferWriteOp::build(OpBuilder &builder, OperationState &result,
Value vector, Value dest, ValueRange indices,
- std::optional<ArrayRef<bool>> inBounds) {
+ std::optional<ArrayRef<bool>> inBounds,
+ /*optional*/ llvm::Align alignment) {
auto vectorType = llvm::cast<VectorType>(vector.getType());
AffineMap permutationMap = getTransferMinorIdentityMap(
llvm::cast<ShapedType>(dest.getType()), vectorType);
- build(builder, result, vector, dest, indices, permutationMap, inBounds);
+ build(builder, result, vector, dest, indices, permutationMap, inBounds,
+ alignment);
}
ParseResult TransferWriteOp::parse(OpAsmParser &parser,
diff --git a/mlir/test/Dialect/Vector/invalid.mlir b/mlir/test/Dialect/Vector/invalid.mlir
index 16295fc23632c..f84e32f8fb647 100644
--- a/mlir/test/Dialect/Vector/invalid.mlir
+++ b/mlir/test/Dialect/Vector/invalid.mlir
@@ -633,6 +633,22 @@ func.func @test_vector.transfer_write(%arg0: memref<?xvector<2xindex>>, %arg1: v
// -----
+func.func @test_vector.transfer_write(%v: vector<1xf32>, %dest: tensor<1xf32>, %cst: index) {
+ // expected-error at +1 {{'vector.transfer_write' 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}}
+ %1 = vector.transfer_write %v, %dest[%cst] {permutation_map = affine_map<(i)->(i)>, alignment = -1} : vector<1xf32>, tensor<1xf32>
+ return
+}
+
+// -----
+
+func.func @test_vector.transfer_write(%v: vector<1xf32>, %dest: tensor<1xf32>, %cst: index) {
+ // expected-error at +1 {{'vector.transfer_write' 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}}
+ %1 = vector.transfer_write %v, %dest[%cst] {permutation_map = affine_map<(i)->(i)>, alignment = 3} : vector<1xf32>, tensor<1xf32>
+ return
+}
+
+// -----
+
func.func @insert_strided_slice(%a: vector<4x4xf32>, %b: vector<4x8x16xf32>) {
// expected-error at +1 {{expected offsets of same size as destination vector rank}}
%1 = vector.insert_strided_slice %a, %b {offsets = [100], strides = [1, 1]} : vector<4x4xf32> into vector<4x8x16xf32>
More information about the Mlir-commits
mailing list