[Mlir-commits] [mlir] 35273dd - [mlir][Vector] Fold chains of `vector.extract` with dynamic positions (#228686)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Fri Oct 23 01:26:05 PDT 2026
Author: Federico Bruzzone
Date: 2026-10-23T10:25:59+02:00
New Revision: 35273ddfdd56b36ce2e73ca0b04e50f1afaaf54b
URL: https://github.com/llvm/llvm-project/commit/35273ddfdd56b36ce2e73ca0b04e50f1afaaf54b
DIFF: https://github.com/llvm/llvm-project/commit/35273ddfdd56b36ce2e73ca0b04e50f1afaaf54b.diff
LOG: [mlir][Vector] Fold chains of `vector.extract` with dynamic positions (#228686)
`foldExtractOpFromExtractChain` bailed out when any extract in the chain
had a dynamic position.
With this PR, we now _simply_ merge the positions as `OpFoldResult`s
instead, then split them back into the `static_position` attribute and
the dynamic operands with `decomposeMixedValues` to correctly set them.
AI-assisted.
---------
Signed-off-by: Federico Bruzzone <federico.bruzzone.i at gmail.com>
Added:
Modified:
mlir/lib/Dialect/Vector/IR/VectorOps.cpp
mlir/test/Conversion/VectorToSCF/vector-to-scf.mlir
mlir/test/Dialect/Vector/canonicalize.mlir
Removed:
################################################################################
diff --git a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
index 490202f2dc9778..eedf483df2d312 100644
--- a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
+++ b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
@@ -1503,32 +1503,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[%idx, 3] : vector<5xf32>
+/// ```
+///
+/// AFTER:
+/// ```mlir
+/// %1 = vector.extract %v[1, 2, %idx, 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 d715ef0fef263b..1de435e8167e5f 100644
--- a/mlir/test/Conversion/VectorToSCF/vector-to-scf.mlir
+++ b/mlir/test/Conversion/VectorToSCF/vector-to-scf.mlir
@@ -879,18 +879,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 10e3184b0a5e0d..4eb3988c417ec7 100644
--- a/mlir/test/Dialect/Vector/canonicalize.mlir
+++ b/mlir/test/Dialect/Vector/canonicalize.mlir
@@ -819,17 +819,26 @@ 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>) {
- %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>
- %d = vector.extract %c[3] : f32 from vector<6xf32>
-
- // CHECK-NEXT: vector.extract %[[A]][0] : vector<4x5x6xf32> from vector<3x4x5x6xf32>
- %e = vector.extract %a[0] : vector<4x5x6xf32> from vector<3x4x5x6xf32>
+// 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, f32, f32) {
+ // Dynamic position (%i) in the inner extract.
+ %b = vector.extract %a[%i] : vector<4x5x6xf32> from vector<3x4x5x6xf32>
+ // CHECK-NEXT: 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.
+ %d = vector.extract %a[0, 1] : vector<5x6xf32> from vector<3x4x5x6xf32>
+ // CHECK-NEXT: 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-NEXT: vector.extract %[[A]][%[[I]], 1, %[[J]], 3] : f32 from vector<3x4x5x6xf32>
+ %g = vector.extract %f[%j, 3] : f32 from vector<5x6xf32>
// CHECK-NEXT: return
- return %d, %e : f32, vector<4x5x6xf32>
+ return %c, %e, %g : f32, f32, f32
}
// -----
@@ -2782,18 +2791,6 @@ func.func @extract_insert_chain(%a: vector<2x16xf32>, %b: vector<12x8x16xf32>, %
// -----
-// CHECK-LABEL: extract_from_extract_chain_should_not_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 {
- %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