[Mlir-commits] [mlir] 501d3fa - [mlir][tensor] Fold extract after insert with dynamic indices (#235765)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Tue Oct 27 06:52:31 PDT 2026
Author: meinie
Date: 2026-10-27T06:52:26-07:00
New Revision: 501d3faadc0a8eced2341051fc3d0b922be6bc1e
URL: https://github.com/llvm/llvm-project/commit/501d3faadc0a8eced2341051fc3d0b922be6bc1e
DIFF: https://github.com/llvm/llvm-project/commit/501d3faadc0a8eced2341051fc3d0b922be6bc1e.diff
LOG: [mlir][tensor] Fold extract after insert with dynamic indices (#235765)
Fold a `tensor.extract` from a `tensor.insert` at the same dynamic
indices to the inserted scalar:
```python
%t1 = tensor.insert %value into %tensor[%i] : tensor<8xf32>
%result = tensor.extract %t1[%i] : tensor<8xf32>
// %result can be replaced with %value.
```
The existing `foldExtractAfterInsert` helper checks matching indices,
but `ExtractOp::fold` returns early while collecting constant indices
when an index is dynamic. Call the helper before that check, leaving its
matching conditions unchanged. This eliminates redundant tensor
construction and extraction without adding a pass or dialect dependency.
Tests cover dynamic and mixed indices, dynamic shapes, nonmatching
indices, shared tensor uses, and repeated insertion.
AI assistance: Code testing and PR description writing are polished with
the assistance of Codex.
Added:
Modified:
mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
mlir/test/Dialect/Tensor/canonicalize.mlir
Removed:
################################################################################
diff --git a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
index 09ce820213dab..9e8bb21aad8f9 100644
--- a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
+++ b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
@@ -1409,6 +1409,9 @@ OpFoldResult ExtractOp::fold(FoldAdaptor adaptor) {
return {};
}
+ if (Value result = foldExtractAfterInsert(*this))
+ return result;
+
// Collect the constant indices into the tensor.
SmallVector<uint64_t, 8> indices;
for (Attribute indice : adaptor.getIndices()) {
@@ -1444,9 +1447,6 @@ OpFoldResult ExtractOp::fold(FoldAdaptor adaptor) {
return elementsAttr.getValues<Attribute>()[indices];
}
- if (Value result = foldExtractAfterInsert(*this))
- return result;
-
return {};
}
diff --git a/mlir/test/Dialect/Tensor/canonicalize.mlir b/mlir/test/Dialect/Tensor/canonicalize.mlir
index abf7004cca5be..690b08620987a 100644
--- a/mlir/test/Dialect/Tensor/canonicalize.mlir
+++ b/mlir/test/Dialect/Tensor/canonicalize.mlir
@@ -226,6 +226,101 @@ func.func @fold_extract(%arg0 : index) -> (f32, f16, f16, i32, complex<f32>, i32
// -----
+// CHECK-LABEL: func.func @extract_after_insert_dynamic(
+// CHECK-SAME: %[[VALUE:.*]]: f32,
+// CHECK-NEXT: return %[[VALUE]] : f32
+func.func @extract_after_insert_dynamic(%value: f32, %tensor: tensor<8xf32>, %i: index) -> f32 {
+ %inserted = tensor.insert %value into %tensor[%i] : tensor<8xf32>
+ %extracted = tensor.extract %inserted[%i] : tensor<8xf32>
+ return %extracted : f32
+}
+
+// -----
+
+// CHECK-LABEL: func.func @extract_after_insert_dynamic_rank2(
+// CHECK-SAME: %[[VALUE:.*]]: i32,
+// CHECK-NEXT: return %[[VALUE]] : i32
+func.func @extract_after_insert_dynamic_rank2(%value: i32, %tensor: tensor<4x8xi32>, %i: index, %j: index) -> i32 {
+ %inserted = tensor.insert %value into %tensor[%i, %j] : tensor<4x8xi32>
+ %extracted = tensor.extract %inserted[%i, %j] : tensor<4x8xi32>
+ return %extracted : i32
+}
+
+// -----
+
+// CHECK-LABEL: func.func @extract_after_insert_mixed_indices(
+// CHECK-SAME: %[[VALUE:.*]]: f32,
+// CHECK-NEXT: return %[[VALUE]] : f32
+func.func @extract_after_insert_mixed_indices(%value: f32, %tensor: tensor<4x8xf32>, %i: index) -> f32 {
+ %c2 = arith.constant 2 : index
+ %inserted = tensor.insert %value into %tensor[%i, %c2] : tensor<4x8xf32>
+ %extracted = tensor.extract %inserted[%i, %c2] : tensor<4x8xf32>
+ return %extracted : f32
+}
+
+// -----
+
+// CHECK-LABEL: func.func @extract_after_insert_dynamic_shape(
+// CHECK-SAME: %[[VALUE:.*]]: f32,
+// CHECK-NEXT: return %[[VALUE]] : f32
+func.func @extract_after_insert_dynamic_shape(%value: f32, %tensor: tensor<?x?xf32>, %i: index, %j: index) -> f32 {
+ %inserted = tensor.insert %value into %tensor[%i, %j] : tensor<?x?xf32>
+ %extracted = tensor.extract %inserted[%i, %j] : tensor<?x?xf32>
+ return %extracted : f32
+}
+
+// -----
+
+// CHECK-LABEL: func.func @extract_after_insert_other_index(
+// CHECK: %[[INSERTED:.*]] = tensor.insert
+// CHECK-NEXT: %[[EXTRACTED:.*]] = tensor.extract %[[INSERTED]]
+// CHECK-NEXT: return %[[EXTRACTED]] : f32
+func.func @extract_after_insert_other_index(%value: f32, %tensor: tensor<8xf32>, %i: index, %j: index) -> f32 {
+ %inserted = tensor.insert %value into %tensor[%i] : tensor<8xf32>
+ %extracted = tensor.extract %inserted[%j] : tensor<8xf32>
+ return %extracted : f32
+}
+
+// -----
+
+// CHECK-LABEL: func.func @extract_after_insert_other_constant(
+// CHECK: %[[INSERTED:.*]] = tensor.insert
+// CHECK-NEXT: %[[EXTRACTED:.*]] = tensor.extract %[[INSERTED]]
+// CHECK-NEXT: return %[[EXTRACTED]] : f32
+func.func @extract_after_insert_other_constant(%value: f32, %tensor: tensor<4x8xf32>, %i: index) -> f32 {
+ %c1 = arith.constant 1 : index
+ %c2 = arith.constant 2 : index
+ %inserted = tensor.insert %value into %tensor[%i, %c1] : tensor<4x8xf32>
+ %extracted = tensor.extract %inserted[%i, %c2] : tensor<4x8xf32>
+ return %extracted : f32
+}
+
+// -----
+
+// CHECK-LABEL: func.func @extract_after_insert_shared_tensor(
+// CHECK-SAME: %[[VALUE:.*]]: f32,
+// CHECK: %[[INSERTED:.*]] = tensor.insert %[[VALUE]]
+// CHECK-NEXT: return %[[INSERTED]], %[[VALUE]] : tensor<8xf32>, f32
+func.func @extract_after_insert_shared_tensor(%value: f32, %tensor: tensor<8xf32>, %i: index) -> (tensor<8xf32>, f32) {
+ %inserted = tensor.insert %value into %tensor[%i] : tensor<8xf32>
+ %extracted = tensor.extract %inserted[%i] : tensor<8xf32>
+ return %inserted, %extracted : tensor<8xf32>, f32
+}
+
+// -----
+
+// CHECK-LABEL: func.func @extract_after_insert_chain(
+// CHECK-SAME: %{{.*}}: f32, %[[VALUE:.*]]: f32,
+// CHECK-NEXT: return %[[VALUE]] : f32
+func.func @extract_after_insert_chain(%first: f32, %value: f32, %tensor: tensor<8xf32>, %i: index) -> f32 {
+ %first_insert = tensor.insert %first into %tensor[%i] : tensor<8xf32>
+ %inserted = tensor.insert %value into %first_insert[%i] : tensor<8xf32>
+ %extracted = tensor.extract %inserted[%i] : tensor<8xf32>
+ return %extracted : f32
+}
+
+// -----
+
// Ensure extract dense resource elements not crash.
// CHECK-LABEL: func @extract_dense_resource_nofold
More information about the Mlir-commits
mailing list