[Mlir-commits] [mlir] [mlir][tensor]-Handle Dynamic Offset in BubbleUpSliceOpThroughCollapse (PR #178921)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Fri Jan 30 09:12:44 PST 2026
llvmbot wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir
@llvm/pr-subscribers-mlir-tensor
Author: Amir Bishara (amirBish)
<details>
<summary>Changes</summary>
This patch extends the `BubbleUpExtractSliceThroughCollapseShape` pattern to handle cases where `tensor.extract_slice` has a dynamic offset.
During tile and fuse transformations, it is common to encounter IR where `tensor.extract_slice` operations appear after `tensor.collapse_shape`. These patterns are used as cleanup transformations to canonicalize the IR by bubbling up the slice operation before the reshape. This enables further optimizations and simplifications downstream.
Previously, the pattern only handled:
1. Static offsets and sizes.
2. Dynamic sizes with a single non-unit expanded dimension.
This left a gap for additional common cases where we may have:
- Dynamic offsets with size == 1 (single element extraction).
- Size greater than 1 but the offset is computed dynamically.
Regarding the first case, It's always legal to perform such a transformation.
Regarding the second case (i.e. dynamic offset with static
size > 1) we can guarantee contiguity by more restricted
conditions:
1. The innermost expanded dimension is divisible by the slice size (ensures the slice fits within a single "row")
2. The offset is provably a multiple of the slice size (ensures the slice starts at an aligned boundary)
For example, given:
- collapse_shape: tensor<6x10xf32> -> tensor<60xf32>
- extract_slice at offset with size
If 10 % size == 0 and offset % 5 == 0, the slice is guaranteed contiguous. The transformation delinearizes the offset to get [row, col] indices and extracts [1, size] from the expanded 6x10 shape.
The offset divisibility check uses affine analysis to statically verify that affine expressions (e.g., `idx * 5`) are multiples of the slice size.
---
Full diff: https://github.com/llvm/llvm-project/pull/178921.diff
2 Files Affected:
- (modified) mlir/lib/Dialect/Tensor/Transforms/ReshapePatterns.cpp (+103-8)
- (modified) mlir/test/Dialect/Tensor/bubble-up-extract-slice-op.mlir (+157-1)
``````````diff
diff --git a/mlir/lib/Dialect/Tensor/Transforms/ReshapePatterns.cpp b/mlir/lib/Dialect/Tensor/Transforms/ReshapePatterns.cpp
index a53af98474245..d069c67daad81 100644
--- a/mlir/lib/Dialect/Tensor/Transforms/ReshapePatterns.cpp
+++ b/mlir/lib/Dialect/Tensor/Transforms/ReshapePatterns.cpp
@@ -579,6 +579,24 @@ LogicalResult mlir::tensor::getCollapsedExtractSliceInfo(
return success();
}
+// Checks if the `value` is a multiple of the `factor`.
+// Otherwise, returns false.
+// For now we are handling only the case where the value
+// is resulted from an affine.apply op, where affine ops
+// can be lowered to.
+static bool isValueMultipleOf(Value value, int64_t factor) {
+ auto applyOp = value.getDefiningOp<affine::AffineApplyOp>();
+ if (!applyOp)
+ return false;
+ AffineMap map = applyOp.getAffineMap();
+ SmallVector<Value> operands(applyOp.getOperands());
+ affine::fullyComposeAffineMapAndOperands(&map, &operands);
+ map = simplifyAffineMap(map);
+ if (map.getNumResults() != 1)
+ return false;
+ return map.getResult(0).isMultipleOf(factor);
+}
+
LogicalResult mlir::tensor::getExpandedExtractSliceInfo(
OpBuilder &b, tensor::ExtractSliceOp sliceOp,
ArrayRef<ReassociationIndices> reassociation,
@@ -610,26 +628,103 @@ LogicalResult mlir::tensor::getExpandedExtractSliceInfo(
for (auto [collapsedSize, collapsedOffset, reassocIndices] :
llvm::zip_equal(collapsedSizes, collapsedOffsets, reassociation)) {
// CASE #1 - size and/or offset are dynamic.
- // In this case, the slice can be represented as a contiguous slice only
- // if there is a single dimension in the reassociation group that has a
- // size not equal to 1.
if (isa<Value>(collapsedSize) || isa<Value>(collapsedOffset)) {
+ // When the size is dynamic, We will handle a simple case where there
+ // is a single dimension in the reassociation group that has a size
+ // not equal to 1, which guarantees it is a contiguous slice.
int nonUnitSizeCount = 0;
+ SmallVector<OpFoldResult> tentativeSizes, tentativeOffsets;
for (int64_t expandedShapeIdx : reassocIndices) {
if (expandedShape[expandedShapeIdx] != 1) {
nonUnitSizeCount++;
- expandedSizes.push_back(collapsedSize);
- expandedOffsets.push_back(collapsedOffset);
+ tentativeSizes.push_back(collapsedSize);
+ tentativeOffsets.push_back(collapsedOffset);
continue;
}
- expandedSizes.push_back(b.getIndexAttr(1));
- expandedOffsets.push_back(b.getIndexAttr(0));
+ tentativeSizes.push_back(b.getIndexAttr(1));
+ tentativeOffsets.push_back(b.getIndexAttr(0));
+ }
+
+ if (nonUnitSizeCount == 1) {
+ // Only one non-unit dimension, slice is contiguous.
+ expandedSizes.append(tentativeSizes);
+ expandedOffsets.append(tentativeOffsets);
+ continue;
}
- if (nonUnitSizeCount != 1) {
+ std::optional<int64_t> staticSize = getConstantIntValue(collapsedSize);
+ if (!staticSize.has_value())
return failure();
+
+ // If the size is statically 1, it means we're extracting a
+ // single element. In this case, we can always handle this by
+ // delinearizing the dynamic offset and get a contiguous slice.
+ if (staticSize.value() == 1) {
+ SmallVector<int64_t> basis;
+ for (int64_t expandedShapeIdx : reassocIndices) {
+ expandedSizes.push_back(b.getIndexAttr(1));
+ basis.push_back(expandedShape[expandedShapeIdx]);
+ }
+
+ Value offsetVal = getValueOrCreateConstantIndexOp(b, sliceOp.getLoc(),
+ collapsedOffset);
+ auto delinearizeOp = affine::AffineDelinearizeIndexOp::create(
+ b, sliceOp.getLoc(), offsetVal, basis, /*hasOuterBound=*/true);
+ for (auto result : delinearizeOp.getResults())
+ expandedOffsets.push_back(result);
+
+ continue;
}
+
+ // If the size is greater than 1, we will take a more
+ // restrictive case where both the offset and the innermost
+ // expanded dimension are divisible by the size.
+ //
+ // Example:
+ // collapse_shape tensor<6x10xf32> -> tensor<60xf32>
+ // extract_slice at offset_a with some size_b.
+ //
+ // Contiguity is guaranteed when:
+ // 1. innermost dim (10) % size_b == 0 i.e.
+ // the slice fits within a single row.
+ // 2. offset_a is multiple of size_b i.e.
+ // the slice starts at row-aligned boundary.
+ //
+ // Result: delinearize offset_a -> [row, col], extract [1, size_b]
+ // from expanded shape.
+ assert(staticSize.value() > 1 && "Expected size to be greater than 1");
+ int64_t sizeVal = staticSize.value();
+ int64_t innermostDim = expandedShape[reassocIndices.back()];
+
+ // Only try this path if offset is from affine.apply
+ Value offsetVal = cast<Value>(collapsedOffset);
+ if (!offsetVal.getDefiningOp<affine::AffineApplyOp>())
+ return failure();
+
+ if (innermostDim % sizeVal != 0)
+ return failure();
+
+ if (!isValueMultipleOf(offsetVal, sizeVal))
+ return failure();
+
+ SmallVector<int64_t> basis;
+ for (int64_t expandedShapeIdx : reassocIndices)
+ basis.push_back(expandedShape[expandedShapeIdx]);
+
+ auto delinearizeOp = affine::AffineDelinearizeIndexOp::create(
+ b, sliceOp.getLoc(), offsetVal, basis, /*hasOuterBound=*/true);
+
+ // Sizes: [1, 1, ..., 1, sizeVal] - size goes on innermost dimension.
+ for (size_t i = 0; i < reassocIndices.size() - 1; ++i)
+ expandedSizes.push_back(b.getIndexAttr(1));
+
+ expandedSizes.push_back(collapsedSize);
+
+ // Offsets from delinearization.
+ for (auto result : delinearizeOp.getResults())
+ expandedOffsets.push_back(result);
+
continue;
}
diff --git a/mlir/test/Dialect/Tensor/bubble-up-extract-slice-op.mlir b/mlir/test/Dialect/Tensor/bubble-up-extract-slice-op.mlir
index 34128d6a5ec8b..97157eb06fb96 100644
--- a/mlir/test/Dialect/Tensor/bubble-up-extract-slice-op.mlir
+++ b/mlir/test/Dialect/Tensor/bubble-up-extract-slice-op.mlir
@@ -190,7 +190,6 @@ func.func @bubble_up_extract_slice_through_collapse_shape_dynamic_size_and_src(%
return %extract : tensor<?xf32>
}
-
// CHECK-LABEL: func.func @bubble_up_extract_slice_through_collapse_shape_dynamic_offset(
// CHECK-SAME: %[[SRC:.*]]: tensor<1x5x1xf32>,
// CHECK-SAME: %[[OFFSET:.*]]: index) -> tensor<3xf32> {
@@ -229,6 +228,163 @@ func.func @bubble_up_extract_slice_through_collapse_shape_dynamic_and_static_gro
return %extract : tensor<20x?xf32>
}
+// CHECK-LABEL: func.func @bubble_up_extract_slice_through_collapse_shape_unit_size_dynamic_offset(
+// CHECK-SAME: %[[SRC:.*]]: tensor<2x3x10xf32>,
+// CHECK-SAME: %[[OFFSET:.*]]: index) -> tensor<1xf32> {
+// CHECK: %[[DELINEARIZE:.*]]:3 = affine.delinearize_index %[[OFFSET]] into (2, 3, 10)
+// CHECK: %[[EXTRACT:.*]] = tensor.extract_slice %[[SRC]]{{\[}}%[[DELINEARIZE]]#0, %[[DELINEARIZE]]#1, %[[DELINEARIZE]]#2] [1, 1, 1] [1, 1, 1]
+// CHECK: %[[COLLAPSE:.*]] = tensor.collapse_shape %[[EXTRACT]] {{\[\[}}0, 1, 2]]
+// CHECK: return %[[COLLAPSE]]
+func.func @bubble_up_extract_slice_through_collapse_shape_unit_size_dynamic_offset(%src: tensor<2x3x10xf32>, %offset : index) -> tensor<1xf32> {
+ %collapse = tensor.collapse_shape %src [[0, 1, 2]] : tensor<2x3x10xf32> into tensor<60xf32>
+ %extract = tensor.extract_slice %collapse[%offset][1][1] : tensor<60xf32> to tensor<1xf32>
+ return %extract : tensor<1xf32>
+}
+
+// CHECK-LABEL: func.func @bubble_up_extract_slice_through_collapse_shape_unit_size_dynamic_offset_multi_group(
+// CHECK-SAME: %[[SRC:.*]]: tensor<2x3x4x5xf32>,
+// CHECK-SAME: %[[OFFSET0:.*]]: index,
+// CHECK-SAME: %[[OFFSET1:.*]]: index) -> tensor<1x1xf32> {
+// CHECK-DAG: %[[DELINEARIZE0:.*]]:2 = affine.delinearize_index %[[OFFSET0]] into (2, 3)
+// CHECK-DAG: %[[DELINEARIZE1:.*]]:2 = affine.delinearize_index %[[OFFSET1]] into (4, 5)
+// CHECK: %[[EXTRACT:.*]] = tensor.extract_slice %[[SRC]]{{\[}}%[[DELINEARIZE0]]#0, %[[DELINEARIZE0]]#1, %[[DELINEARIZE1]]#0, %[[DELINEARIZE1]]#1] [1, 1, 1, 1] [1, 1, 1, 1]
+// CHECK: %[[COLLAPSE:.*]] = tensor.collapse_shape %[[EXTRACT]] {{\[\[}}0, 1], [2, 3]]
+// CHECK: return %[[COLLAPSE]]
+func.func @bubble_up_extract_slice_through_collapse_shape_unit_size_dynamic_offset_multi_group(%src: tensor<2x3x4x5xf32>, %offset0 : index, %offset1 : index) -> tensor<1x1xf32> {
+ %collapse = tensor.collapse_shape %src [[0, 1], [2, 3]] : tensor<2x3x4x5xf32> into tensor<6x20xf32>
+ %extract = tensor.extract_slice %collapse[%offset0, %offset1][1, 1][1, 1] : tensor<6x20xf32> to tensor<1x1xf32>
+ return %extract : tensor<1x1xf32>
+}
+
+// CHECK-LABEL: func.func @bubble_up_extract_slice_through_collapse_shape_affine_apply_offset(
+// CHECK-SAME: %[[SRC:.*]]: tensor<2x3x10xf32>,
+// CHECK-SAME: %[[IDX:.*]]: index) -> tensor<5xf32> {
+// CHECK: %[[OFFSET:.*]] = affine.apply
+// CHECK: %[[DELINEARIZE:.*]]:3 = affine.delinearize_index %[[OFFSET]] into (2, 3, 10)
+// CHECK: %[[EXTRACT:.*]] = tensor.extract_slice %[[SRC]]{{\[}}%[[DELINEARIZE]]#0, %[[DELINEARIZE]]#1, %[[DELINEARIZE]]#2] [1, 1, 5] [1, 1, 1]
+// CHECK: %[[COLLAPSE:.*]] = tensor.collapse_shape %[[EXTRACT]] {{\[\[}}0, 1, 2]]
+// CHECK: return %[[COLLAPSE]]
+func.func @bubble_up_extract_slice_through_collapse_shape_affine_apply_offset(%src: tensor<2x3x10xf32>, %idx : index) -> tensor<5xf32> {
+ %collapse = tensor.collapse_shape %src [[0, 1, 2]] : tensor<2x3x10xf32> into tensor<60xf32>
+ %offset = affine.apply affine_map<()[s0] -> (s0 * 5)>()[%idx]
+ %extract = tensor.extract_slice %collapse[%offset][5][1] : tensor<60xf32> to tensor<5xf32>
+ return %extract : tensor<5xf32>
+}
+
+// CHECK-LABEL: func.func @bubble_up_extract_slice_through_collapse_shape_affine_apply_offset_full_inner_dim(
+// CHECK-SAME: %[[SRC:.*]]: tensor<4x5x8xf32>,
+// CHECK-SAME: %[[IDX:.*]]: index) -> tensor<8xf32> {
+// CHECK: %[[OFFSET:.*]] = affine.apply
+// CHECK: %[[DELINEARIZE:.*]]:3 = affine.delinearize_index %[[OFFSET]] into (4, 5, 8)
+// CHECK: %[[EXTRACT:.*]] = tensor.extract_slice %[[SRC]]{{\[}}%[[DELINEARIZE]]#0, %[[DELINEARIZE]]#1, %[[DELINEARIZE]]#2] [1, 1, 8] [1, 1, 1]
+// CHECK: %[[COLLAPSE:.*]] = tensor.collapse_shape %[[EXTRACT]] {{\[\[}}0, 1, 2]]
+// CHECK: return %[[COLLAPSE]]
+func.func @bubble_up_extract_slice_through_collapse_shape_affine_apply_offset_full_inner_dim(%src: tensor<4x5x8xf32>, %idx : index) -> tensor<8xf32> {
+ %collapse = tensor.collapse_shape %src [[0, 1, 2]] : tensor<4x5x8xf32> into tensor<160xf32>
+ // offset = idx * 8, size = 8, innermost dim = 8, 8 % 8 == 0, should work
+ %offset = affine.apply affine_map<()[s0] -> (s0 * 8)>()[%idx]
+ %extract = tensor.extract_slice %collapse[%offset][8][1] : tensor<160xf32> to tensor<8xf32>
+ return %extract : tensor<8xf32>
+}
+
+// CHECK-LABEL: func.func @bubble_up_extract_slice_through_collapse_shape_affine_apply_complex_expr(
+// CHECK-SAME: %[[SRC:.*]]: tensor<3x4x8xf32>,
+// CHECK-SAME: %[[IDX:.*]]: index) -> tensor<4xf32> {
+// CHECK: %[[OFFSET:.*]] = affine.apply
+// CHECK: %[[DELINEARIZE:.*]]:3 = affine.delinearize_index %[[OFFSET]] into (3, 4, 8)
+// CHECK: %[[EXTRACT:.*]] = tensor.extract_slice %[[SRC]]{{\[}}%[[DELINEARIZE]]#0, %[[DELINEARIZE]]#1, %[[DELINEARIZE]]#2] [1, 1, 4] [1, 1, 1]
+// CHECK: %[[COLLAPSE:.*]] = tensor.collapse_shape %[[EXTRACT]] {{\[\[}}0, 1, 2]]
+// CHECK: return %[[COLLAPSE]]
+func.func @bubble_up_extract_slice_through_collapse_shape_affine_apply_complex_expr(%src: tensor<3x4x8xf32>, %idx : index) -> tensor<4xf32> {
+ %collapse = tensor.collapse_shape %src [[0, 1, 2]] : tensor<3x4x8xf32> into tensor<96xf32>
+ // offset = idx * 4, size = 4, innermost dim = 8, 8 % 4 == 0, should work
+ %offset = affine.apply affine_map<()[s0] -> (s0 * 4)>()[%idx]
+ %extract = tensor.extract_slice %collapse[%offset][4][1] : tensor<96xf32> to tensor<4xf32>
+ return %extract : tensor<4xf32>
+}
+
+// CHECK-LABEL: func.func @bubble_up_extract_slice_through_collapse_shape_affine_apply_multi_group(
+// CHECK-SAME: %[[SRC:.*]]: tensor<2x3x4x6xf32>,
+// CHECK-SAME: %[[IDX0:.*]]: index,
+// CHECK-SAME: %[[IDX1:.*]]: index) -> tensor<3x2xf32> {
+// CHECK-DAG: %[[OFFSET0:.*]] = affine.apply {{.*}}{{\[}}%[[IDX0]]]
+// CHECK-DAG: %[[OFFSET1:.*]] = affine.apply {{.*}}{{\[}}%[[IDX1]]]
+// CHECK-DAG: %[[DELINEARIZE0:.*]]:2 = affine.delinearize_index %[[OFFSET0]] into (2, 3)
+// CHECK-DAG: %[[DELINEARIZE1:.*]]:2 = affine.delinearize_index %[[OFFSET1]] into (4, 6)
+// CHECK: %[[EXTRACT:.*]] = tensor.extract_slice %[[SRC]]{{\[}}%[[DELINEARIZE0]]#0, %[[DELINEARIZE0]]#1, %[[DELINEARIZE1]]#0, %[[DELINEARIZE1]]#1] [1, 3, 1, 2] [1, 1, 1, 1]
+// CHECK: %[[COLLAPSE:.*]] = tensor.collapse_shape %[[EXTRACT]] {{\[\[}}0, 1], [2, 3]]
+// CHECK: return %[[COLLAPSE]]
+func.func @bubble_up_extract_slice_through_collapse_shape_affine_apply_multi_group(%src: tensor<2x3x4x6xf32>, %idx0 : index, %idx1 : index) -> tensor<3x2xf32> {
+ %collapse = tensor.collapse_shape %src [[0, 1], [2, 3]] : tensor<2x3x4x6xf32> into tensor<6x24xf32>
+ // For first group: size=3, innermost=3, 3%3==0, offset=idx0*3, works
+ // For second group: size=2, innermost=6, 6%2==0, offset=idx1*2, works
+ %offset0 = affine.apply affine_map<()[s0] -> (s0 * 3)>()[%idx0]
+ %offset1 = affine.apply affine_map<()[s0] -> (s0 * 2)>()[%idx1]
+ %extract = tensor.extract_slice %collapse[%offset0, %offset1][3, 2][1, 1] : tensor<6x24xf32> to tensor<3x2xf32>
+ return %extract : tensor<3x2xf32>
+}
+
+// CHECK-LABEL: func.func @bubble_up_extract_slice_through_collapse_shape_affine_apply_nested(
+// CHECK-SAME: %[[SRC:.*]]: tensor<2x4x8xf32>,
+// CHECK-SAME: %[[IDX:.*]]: index) -> tensor<4xf32> {
+// CHECK: affine.apply
+// CHECK: %[[OFFSET:.*]] = affine.apply
+// CHECK: %[[DELINEARIZE:.*]]:3 = affine.delinearize_index %[[OFFSET]] into (2, 4, 8)
+// CHECK: %[[EXTRACT:.*]] = tensor.extract_slice %[[SRC]]{{\[}}%[[DELINEARIZE]]#0, %[[DELINEARIZE]]#1, %[[DELINEARIZE]]#2] [1, 1, 4] [1, 1, 1]
+// CHECK: %[[COLLAPSE:.*]] = tensor.collapse_shape %[[EXTRACT]] {{\[\[}}0, 1, 2]]
+// CHECK: return %[[COLLAPSE]]
+func.func @bubble_up_extract_slice_through_collapse_shape_affine_apply_nested(%src: tensor<2x4x8xf32>, %idx : index) -> tensor<4xf32> {
+ %collapse = tensor.collapse_shape %src [[0, 1, 2]] : tensor<2x4x8xf32> into tensor<64xf32>
+ // Nested affine.apply: offset = (idx * 2) * 2 = idx * 4, multiple of 4
+ %tmp = affine.apply affine_map<()[s0] -> (s0 * 2)>()[%idx]
+ %offset = affine.apply affine_map<()[s0] -> (s0 * 2)>()[%tmp]
+ %extract = tensor.extract_slice %collapse[%offset][4][1] : tensor<64xf32> to tensor<4xf32>
+ return %extract : tensor<4xf32>
+}
+
+// CHECK-LABEL: func.func @no_bubble_up_extract_slice_affine_apply_offset_not_multiple(
+// CHECK-SAME: %[[SRC:.*]]: tensor<2x3x10xf32>,
+// CHECK-SAME: %[[IDX:.*]]: index) -> tensor<4xf32> {
+// CHECK: %[[COLLAPSE:.*]] = tensor.collapse_shape
+// CHECK: %[[OFFSET:.*]] = affine.apply
+// CHECK: %[[EXTRACT:.*]] = tensor.extract_slice
+// CHECK: return %[[EXTRACT]]
+func.func @no_bubble_up_extract_slice_affine_apply_offset_not_multiple(%src: tensor<2x3x10xf32>, %idx : index) -> tensor<4xf32> {
+ %collapse = tensor.collapse_shape %src [[0, 1, 2]] : tensor<2x3x10xf32> into tensor<60xf32>
+ // offset = idx * 3, but size = 4, 3 is not a multiple of 4
+ %offset = affine.apply affine_map<()[s0] -> (s0 * 3)>()[%idx]
+ %extract = tensor.extract_slice %collapse[%offset][4][1] : tensor<60xf32> to tensor<4xf32>
+ return %extract : tensor<4xf32>
+}
+
+// CHECK-LABEL: func.func @no_bubble_up_extract_slice_affine_apply_size_not_dividing_inner(
+// CHECK-SAME: %[[SRC:.*]]: tensor<2x3x10xf32>,
+// CHECK-SAME: %[[IDX:.*]]: index) -> tensor<3xf32> {
+// CHECK: %[[COLLAPSE:.*]] = tensor.collapse_shape
+// CHECK: %[[OFFSET:.*]] = affine.apply
+// CHECK: %[[EXTRACT:.*]] = tensor.extract_slice
+// CHECK: return %[[EXTRACT]]
+func.func @no_bubble_up_extract_slice_affine_apply_size_not_dividing_inner(%src: tensor<2x3x10xf32>, %idx : index) -> tensor<3xf32> {
+ %collapse = tensor.collapse_shape %src [[0, 1, 2]] : tensor<2x3x10xf32> into tensor<60xf32>
+ // offset = idx * 3 (multiple of 3), but size = 3, innermost = 10, 10 % 3 != 0
+ %offset = affine.apply affine_map<()[s0] -> (s0 * 3)>()[%idx]
+ %extract = tensor.extract_slice %collapse[%offset][3][1] : tensor<60xf32> to tensor<3xf32>
+ return %extract : tensor<3xf32>
+}
+
+// CHECK-LABEL: func.func @no_bubble_up_extract_slice_dynamic_offset_not_affine_apply(
+// CHECK-SAME: %[[SRC:.*]]: tensor<2x3x10xf32>,
+// CHECK-SAME: %[[OFFSET:.*]]: index) -> tensor<5xf32> {
+// CHECK: %[[COLLAPSE:.*]] = tensor.collapse_shape
+// CHECK: %[[EXTRACT:.*]] = tensor.extract_slice
+// CHECK: return %[[EXTRACT]]
+func.func @no_bubble_up_extract_slice_dynamic_offset_not_affine_apply(%src: tensor<2x3x10xf32>, %offset : index) -> tensor<5xf32> {
+ %collapse = tensor.collapse_shape %src [[0, 1, 2]] : tensor<2x3x10xf32> into tensor<60xf32>
+ // offset is a plain Value (not from affine.apply), can't prove it's a multiple of 5
+ %extract = tensor.extract_slice %collapse[%offset][5][1] : tensor<60xf32> to tensor<5xf32>
+ return %extract : tensor<5xf32>
+}
+
/// The 2 following tests are cases where the bubble up cannot occur because the contiguous size extracted
/// from the collapsed shape cannot be expressed via a single extract_slice op.
/// In the first test it is because the size extracted cannot be expressed as a slice
``````````
</details>
https://github.com/llvm/llvm-project/pull/178921
More information about the Mlir-commits
mailing list