[Mlir-commits] [mlir] [mlir][Vector] Fold chains of `vector.extract` with dynamic positions (PR #228686)
Federico Bruzzone
llvmlistbot at llvm.org
Sun Oct 4 15:52:37 PDT 2026
https://github.com/FedericoBruzzone updated https://github.com/llvm/llvm-project/pull/228686
>From e19a0a81cb49f457132de299e5dbf39d199e17e3 Mon Sep 17 00:00:00 2001
From: Federico Bruzzone <federico.bruzzone.i at gmail.com>
Date: Sat, 3 Oct 2026 11:56:19 +0200
Subject: [PATCH 1/2] [mlir][Vector] Fold chains of `vector.extract` with
dynamic positions
Signed-off-by: Federico Bruzzone <federico.bruzzone.i at gmail.com>
---
mlir/lib/Dialect/Vector/IR/VectorOps.cpp | 38 ++++++++++---------
.../Conversion/VectorToSCF/vector-to-scf.mlir | 12 ++----
mlir/test/Dialect/Vector/canonicalize.mlir | 32 ++++++++++++++--
3 files changed, 52 insertions(+), 30 deletions(-)
diff --git a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
index 638f7017bc581..b48ef2b9830e5 100644
--- a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
+++ b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
@@ -1482,32 +1482,34 @@ static SmallVector<IntType> extractVector(ArrayAttr arrayAttr) {
}
/// Fold the result of chains of ExtractOp in place by simply concatenating the
-/// positions.
+/// positions. Both static and dynamic positions are supported.
+///
+/// BEFORE:
+/// ```mlir
+/// %0 = vector.extract %v[1, 2] : vector<3x4x5xf32>
+/// %1 = vector.extract %0[%c0, 3] : vector<5xf32>
+/// ```
+///
+/// AFTER:
+/// ```mlir
+/// %1 = vector.extract %v[1, 2, %c0, 3] : vector<3x4x5xf32>
+/// ```
static LogicalResult foldExtractOpFromExtractChain(ExtractOp extractOp) {
if (!extractOp.getSource().getDefiningOp<ExtractOp>())
return failure();
- // TODO: Canonicalization for dynamic position not implemented yet.
- if (extractOp.hasDynamicPosition())
- return failure();
-
- SmallVector<int64_t> globalPosition;
+ SmallVector<OpFoldResult> globalPosition = extractOp.getMixedPosition();
ExtractOp currentOp = extractOp;
- ArrayRef<int64_t> extrPos = currentOp.getStaticPosition();
- globalPosition.append(extrPos.rbegin(), extrPos.rend());
while (ExtractOp nextOp = currentOp.getSource().getDefiningOp<ExtractOp>()) {
currentOp = nextOp;
- // TODO: Canonicalization for dynamic position not implemented yet.
- if (currentOp.hasDynamicPosition())
- return failure();
- ArrayRef<int64_t> extrPos = currentOp.getStaticPosition();
- globalPosition.append(extrPos.rbegin(), extrPos.rend());
+ SmallVector<OpFoldResult> pos = currentOp.getMixedPosition();
+ globalPosition.insert(globalPosition.begin(), pos.begin(), pos.end());
}
- extractOp.setOperand(0, currentOp.getSource());
- // OpBuilder is only used as a helper to build an I64ArrayAttr.
- OpBuilder b(extractOp.getContext());
- std::reverse(globalPosition.begin(), globalPosition.end());
- extractOp.setStaticPosition(globalPosition);
+ auto [staticPos, dynPos] = decomposeMixedValues(globalPosition);
+ Value source = currentOp.getSource();
+ extractOp->setOperands(
+ llvm::to_vector(llvm::concat<Value>(ValueRange(source), dynPos)));
+ extractOp.setStaticPosition(staticPos);
return success();
}
diff --git a/mlir/test/Conversion/VectorToSCF/vector-to-scf.mlir b/mlir/test/Conversion/VectorToSCF/vector-to-scf.mlir
index e3a7f57e144ed..0a35546408a63 100644
--- a/mlir/test/Conversion/VectorToSCF/vector-to-scf.mlir
+++ b/mlir/test/Conversion/VectorToSCF/vector-to-scf.mlir
@@ -819,18 +819,14 @@ func.func @scalable_transpose_store_unmasked(%vec: vector<4x[4]xf32>, %dest: mem
// FULL-UNROLL-DAG: %[[C0:.*]] = arith.constant 0 : index
// FULL-UNROLL-DAG: %[[C1:.*]] = arith.constant 1 : index
// FULL-UNROLL-DAG: %[[C4:.*]] = arith.constant 4 : index
-// FULL-UNROLL: %[[SLICE_0:.*]] = vector.extract %[[VEC]][0] : vector<[4]xf32> from vector<4x[4]xf32>
-// FULL-UNROLL: %[[SLICE_1:.*]] = vector.extract %[[VEC]][1] : vector<[4]xf32> from vector<4x[4]xf32>
-// FULL-UNROLL: %[[SLICE_2:.*]] = vector.extract %[[VEC]][2] : vector<[4]xf32> from vector<4x[4]xf32>
-// FULL-UNROLL: %[[SLICE_3:.*]] = vector.extract %[[VEC]][3] : vector<[4]xf32> from vector<4x[4]xf32>
// FULL-UNROLL: %[[VSCALE:.*]] = vector.vscale
// FULL-UNROLL: %[[C4_VSCALE:.*]] = arith.muli %[[VSCALE]], %[[C4]] : index
// FULL-UNROLL: scf.for %[[VAL_13:.*]] = %[[C0]] to %[[C4_VSCALE]] step %[[C1]] {
// FULL-UNROLL: %[[SLICE_I:.*]] = affine.apply #[[$SLICE_MAP]](%[[VAL_13]]){{\[}}%[[I]]]
-// FULL-UNROLL: %[[ELEM_0:.*]] = vector.extract %[[SLICE_0]]{{\[}}%[[VAL_13]]] : f32 from vector<[4]xf32>
-// FULL-UNROLL: %[[ELEM_1:.*]] = vector.extract %[[SLICE_1]]{{\[}}%[[VAL_13]]] : f32 from vector<[4]xf32>
-// FULL-UNROLL: %[[ELEM_2:.*]] = vector.extract %[[SLICE_2]]{{\[}}%[[VAL_13]]] : f32 from vector<[4]xf32>
-// FULL-UNROLL: %[[ELEM_3:.*]] = vector.extract %[[SLICE_3]]{{\[}}%[[VAL_13]]] : f32 from vector<[4]xf32>
+// FULL-UNROLL: %[[ELEM_0:.*]] = vector.extract %[[VEC]][0, %[[VAL_13]]] : f32 from vector<4x[4]xf32>
+// FULL-UNROLL: %[[ELEM_1:.*]] = vector.extract %[[VEC]][1, %[[VAL_13]]] : f32 from vector<4x[4]xf32>
+// FULL-UNROLL: %[[ELEM_2:.*]] = vector.extract %[[VEC]][2, %[[VAL_13]]] : f32 from vector<4x[4]xf32>
+// FULL-UNROLL: %[[ELEM_3:.*]] = vector.extract %[[VEC]][3, %[[VAL_13]]] : f32 from vector<4x[4]xf32>
// FULL-UNROLL: %[[TRANSPOSE_SLICE:.*]] = vector.from_elements %[[ELEM_0]], %[[ELEM_1]], %[[ELEM_2]], %[[ELEM_3]] : vector<4xf32>
// FULL-UNROLL: vector.transfer_write %[[TRANSPOSE_SLICE]], %[[DEST]]{{\[}}%[[SLICE_I]], %[[J]]] {in_bounds = [true]} : vector<4xf32>, memref<?x?xf32>
diff --git a/mlir/test/Dialect/Vector/canonicalize.mlir b/mlir/test/Dialect/Vector/canonicalize.mlir
index 6b271ee5cce19..a55dd11b45b3d 100644
--- a/mlir/test/Dialect/Vector/canonicalize.mlir
+++ b/mlir/test/Dialect/Vector/canonicalize.mlir
@@ -834,6 +834,31 @@ func.func @fold_extracts(%a : vector<3x4x5x6xf32>) -> (f32, vector<4x5x6xf32>) {
// -----
+// CHECK-LABEL: fold_extracts_dynamic_position
+// CHECK-SAME: %[[A:[a-zA-Z0-9]*]]: vector<3x4x5x6xf32>
+// CHECK-SAME: %[[I:[a-zA-Z0-9]*]]: index
+// CHECK-SAME: %[[J:[a-zA-Z0-9]*]]: index
+func.func @fold_extracts_dynamic_position(%a : vector<3x4x5x6xf32>, %i : index, %j : index) -> (f32, f32, f32) {
+ // Dynamic position (%i) in the inner extract (%b).
+ %b = vector.extract %a[%i] : vector<4x5x6xf32> from vector<3x4x5x6xf32>
+ // CHECK: vector.extract %[[A]][%[[I]], 1, 2, 3] : f32 from vector<3x4x5x6xf32>
+ %c = vector.extract %b[1, 2, 3] : f32 from vector<4x5x6xf32>
+
+ // Dynamic position (%j) in the outer extract (%e).
+ %d = vector.extract %a[0, 1] : vector<5x6xf32> from vector<3x4x5x6xf32>
+ // CHECK: vector.extract %[[A]][0, 1, %[[J]], 3] : f32 from vector<3x4x5x6xf32>
+ %e = vector.extract %d[%j, 3] : f32 from vector<5x6xf32>
+
+ // Dynamic positions (%i, %j) in both extracts.
+ %f = vector.extract %a[%i, 1] : vector<5x6xf32> from vector<3x4x5x6xf32>
+ // CHECK: vector.extract %[[A]][%[[I]], 1, %[[J]], 3] : f32 from vector<3x4x5x6xf32>
+ %g = vector.extract %f[%j, 3] : f32 from vector<5x6xf32>
+
+ return %c, %e, %g : f32, f32, f32
+}
+
+// -----
+
// CHECK-LABEL: fold_extract_transpose
// CHECK-SAME: %[[A:[a-zA-Z0-9]*]]: vector<3x4x5x6xf32>
// CHECK-SAME: %[[B:[a-zA-Z0-9]*]]: vector<3x6x5x6xf32>
@@ -2710,11 +2735,10 @@ func.func @extract_insert_chain(%a: vector<2x16xf32>, %b: vector<12x8x16xf32>, %
// -----
-// CHECK-LABEL: extract_from_extract_chain_should_not_fold_dynamic_extracts
+// CHECK-LABEL: extract_from_extract_chain_fold_dynamic_extracts
// CHECK-SAME: (%[[VEC:.*]]: vector<2x4xf32>, %[[IDX:.*]]: index)
-// CHECK: %[[A:.*]] = vector.extract %[[VEC]][%[[IDX]]] : vector<4xf32> from vector<2x4xf32>
-// CHECK: %[[B:.*]] = vector.extract %[[A]][1] : f32 from vector<4xf32>
-func.func @extract_from_extract_chain_should_not_fold_dynamic_extracts(%v: vector<2x4xf32>, %index: index) -> f32 {
+// CHECK: vector.extract %[[VEC]][%[[IDX]], 1] : f32 from vector<2x4xf32>
+func.func @extract_from_extract_chain_fold_dynamic_extracts(%v: vector<2x4xf32>, %index: index) -> f32 {
%0 = vector.extract %v[%index] : vector<4xf32> from vector<2x4xf32>
%1 = vector.extract %0[1] : f32 from vector<4xf32>
return %1 : f32
>From a6ee5d77bf05600e1de4ebe0c0639fd1b86ea049 Mon Sep 17 00:00:00 2001
From: Federico Bruzzone <federico.bruzzone.i at gmail.com>
Date: Mon, 5 Oct 2026 00:51:56 +0200
Subject: [PATCH 2/2] Address comments from Andrzej
Signed-off-by: Federico Bruzzone <federico.bruzzone.i at gmail.com>
---
mlir/lib/Dialect/Vector/IR/VectorOps.cpp | 4 +-
mlir/test/Dialect/Vector/canonicalize.mlir | 52 +++++++---------------
2 files changed, 19 insertions(+), 37 deletions(-)
diff --git a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
index b48ef2b9830e5..681072b38662b 100644
--- a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
+++ b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
@@ -1487,12 +1487,12 @@ static SmallVector<IntType> extractVector(ArrayAttr arrayAttr) {
/// BEFORE:
/// ```mlir
/// %0 = vector.extract %v[1, 2] : vector<3x4x5xf32>
-/// %1 = vector.extract %0[%c0, 3] : vector<5xf32>
+/// %1 = vector.extract %0[%idx, 3] : vector<5xf32>
/// ```
///
/// AFTER:
/// ```mlir
-/// %1 = vector.extract %v[1, 2, %c0, 3] : vector<3x4x5xf32>
+/// %1 = vector.extract %v[1, 2, %idx, 3] : vector<3x4x5xf32>
/// ```
static LogicalResult foldExtractOpFromExtractChain(ExtractOp extractOp) {
if (!extractOp.getSource().getDefiningOp<ExtractOp>())
diff --git a/mlir/test/Dialect/Vector/canonicalize.mlir b/mlir/test/Dialect/Vector/canonicalize.mlir
index a55dd11b45b3d..3692f9f375642 100644
--- a/mlir/test/Dialect/Vector/canonicalize.mlir
+++ b/mlir/test/Dialect/Vector/canonicalize.mlir
@@ -819,7 +819,10 @@ func.func @insert_extract_transpose_3d(
// CHECK-LABEL: fold_extracts
// CHECK-SAME: %[[A:[a-zA-Z0-9]*]]: vector<3x4x5x6xf32>
-func.func @fold_extracts(%a : vector<3x4x5x6xf32>) -> (f32, vector<4x5x6xf32>) {
+// CHECK-SAME: %[[I:[a-zA-Z0-9]*]]: index
+// CHECK-SAME: %[[J:[a-zA-Z0-9]*]]: index
+func.func @fold_extracts(%a : vector<3x4x5x6xf32>, %i : index, %j : index) -> (f32, vector<4x5x6xf32>, f32, f32, f32) {
+ // Static positions only.
%b = vector.extract %a[0] : vector<4x5x6xf32> from vector<3x4x5x6xf32>
%c = vector.extract %b[1, 2] : vector<6xf32> from vector<4x5x6xf32>
// CHECK-NEXT: vector.extract %[[A]][0, 1, 2, 3] : f32 from vector<3x4x5x6xf32>
@@ -828,33 +831,23 @@ func.func @fold_extracts(%a : vector<3x4x5x6xf32>) -> (f32, vector<4x5x6xf32>) {
// CHECK-NEXT: vector.extract %[[A]][0] : vector<4x5x6xf32> from vector<3x4x5x6xf32>
%e = vector.extract %a[0] : vector<4x5x6xf32> from vector<3x4x5x6xf32>
- // CHECK-NEXT: return
- return %d, %e : f32, vector<4x5x6xf32>
-}
-
-// -----
-
-// CHECK-LABEL: fold_extracts_dynamic_position
-// CHECK-SAME: %[[A:[a-zA-Z0-9]*]]: vector<3x4x5x6xf32>
-// CHECK-SAME: %[[I:[a-zA-Z0-9]*]]: index
-// CHECK-SAME: %[[J:[a-zA-Z0-9]*]]: index
-func.func @fold_extracts_dynamic_position(%a : vector<3x4x5x6xf32>, %i : index, %j : index) -> (f32, f32, f32) {
- // Dynamic position (%i) in the inner extract (%b).
- %b = vector.extract %a[%i] : vector<4x5x6xf32> from vector<3x4x5x6xf32>
- // CHECK: vector.extract %[[A]][%[[I]], 1, 2, 3] : f32 from vector<3x4x5x6xf32>
- %c = vector.extract %b[1, 2, 3] : f32 from vector<4x5x6xf32>
+ // Dynamic position (%i) in the inner extract.
+ %f = vector.extract %a[%i] : vector<4x5x6xf32> from vector<3x4x5x6xf32>
+ // CHECK-NEXT: vector.extract %[[A]][%[[I]], 1, 2, 3] : f32 from vector<3x4x5x6xf32>
+ %g = vector.extract %f[1, 2, 3] : f32 from vector<4x5x6xf32>
- // Dynamic position (%j) in the outer extract (%e).
- %d = vector.extract %a[0, 1] : vector<5x6xf32> from vector<3x4x5x6xf32>
- // CHECK: vector.extract %[[A]][0, 1, %[[J]], 3] : f32 from vector<3x4x5x6xf32>
- %e = vector.extract %d[%j, 3] : f32 from vector<5x6xf32>
+ // Dynamic position (%j) in the outer extract.
+ %h = vector.extract %a[0, 1] : vector<5x6xf32> from vector<3x4x5x6xf32>
+ // CHECK-NEXT: vector.extract %[[A]][0, 1, %[[J]], 3] : f32 from vector<3x4x5x6xf32>
+ %k = vector.extract %h[%j, 3] : f32 from vector<5x6xf32>
// Dynamic positions (%i, %j) in both extracts.
- %f = vector.extract %a[%i, 1] : vector<5x6xf32> from vector<3x4x5x6xf32>
- // CHECK: vector.extract %[[A]][%[[I]], 1, %[[J]], 3] : f32 from vector<3x4x5x6xf32>
- %g = vector.extract %f[%j, 3] : f32 from vector<5x6xf32>
+ %l = vector.extract %a[%i, 1] : vector<5x6xf32> from vector<3x4x5x6xf32>
+ // CHECK-NEXT: vector.extract %[[A]][%[[I]], 1, %[[J]], 3] : f32 from vector<3x4x5x6xf32>
+ %m = vector.extract %l[%j, 3] : f32 from vector<5x6xf32>
- return %c, %e, %g : f32, f32, f32
+ // CHECK-NEXT: return
+ return %d, %e, %g, %k, %m : f32, vector<4x5x6xf32>, f32, f32, f32
}
// -----
@@ -2735,17 +2728,6 @@ func.func @extract_insert_chain(%a: vector<2x16xf32>, %b: vector<12x8x16xf32>, %
// -----
-// CHECK-LABEL: extract_from_extract_chain_fold_dynamic_extracts
-// CHECK-SAME: (%[[VEC:.*]]: vector<2x4xf32>, %[[IDX:.*]]: index)
-// CHECK: vector.extract %[[VEC]][%[[IDX]], 1] : f32 from vector<2x4xf32>
-func.func @extract_from_extract_chain_fold_dynamic_extracts(%v: vector<2x4xf32>, %index: index) -> f32 {
- %0 = vector.extract %v[%index] : vector<4xf32> from vector<2x4xf32>
- %1 = vector.extract %0[1] : f32 from vector<4xf32>
- return %1 : f32
-}
-
-// -----
-
// CHECK-LABEL: extract_extract_strided2
// CHECK-SAME: %[[A:.*]]: vector<2x4xf32>
// CHECK: %[[V:.*]] = vector.extract %[[A]][1] : vector<4xf32> from vector<2x4xf32>
More information about the Mlir-commits
mailing list