[MLIR] Bubble up tensor.extract_slice through tensor.collapse_shape #131982
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Add a pattern that bubbles up tensor.extract_slice through tensor.collapse_shape. The pattern is registered in a pattern population function that is used by the transform op transform.apply_patterns.tensor.bubble_up_extract_slice and by the tranform op transform.structured.fuse as a cleanup pattern. This pattern enables tiling and fusing op chains which contain tensor.collapse_shape if added as a cleanup pattern of tile and fuse utility. Without this pattern that would not be possible, as tensor.collapse_shape does not implement the tiling interface. This is an additional pattern to the one added in PR llvm#126898
ofri-frishman
requested review from
hanhanW and
nicolasvasilache
as code owners
|
@llvm/pr-subscribers-mlir @llvm/pr-subscribers-mlir-tensor Author: ofri frishman (ofri-frishman) ChangesAdd a pattern that bubbles up tensor.extract_slice through tensor.collapse_shape. Patch is 24.24 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/131982.diff 3 Files Affected:
diff --git a/mlir/lib/Dialect/Tensor/Transforms/ReshapePatterns.cpp b/mlir/lib/Dialect/Tensor/Transforms/ReshapePatterns.cpp
index acedf51d0e240..efa4d10817e39 100644
--- a/mlir/lib/Dialect/Tensor/Transforms/ReshapePatterns.cpp
+++ b/mlir/lib/Dialect/Tensor/Transforms/ReshapePatterns.cpp
@@ -12,8 +12,10 @@
#include "mlir/Dialect/Tensor/Transforms/Transforms.h"
#include "mlir/IR/PatternMatch.h"
#include "mlir/Interfaces/ValueBoundsOpInterface.h"
+#include "llvm/ADT/STLExtras.h"
#include "llvm/Support/Debug.h"
#include "llvm/Support/LogicalResult.h"
+#include <algorithm>
using namespace mlir;
using namespace mlir::tensor;
@@ -428,6 +430,190 @@ struct BubbleUpExpandShapeThroughExtractSlice
}
};
+/// Converts `tensor.collapse_shape(tensor.extract_slice)` to
+/// `tensor.extract_slice(tensor.collapse_shape)`.
+///
+/// For this transformation to be possible, the slice must be representable as a
+/// contiguous slice within each reassociation group of the src.
+///
+/// In case the size and offset extracted are static then this is possible if
+/// the following conditions are met:
+/// Let T be a tensor of shape [A0, A1, ..., An], and let S = [S0, S1, ..., Sn]
+/// be the shape of a desired slice. A slice of shape S can be extracted as a
+/// contiguous block of memory if and only if there exists an index k in {0, 1,
+/// ..., n} such that:
+/// S_i = 1 for all i < k (that is, all leading dimensions are singleton),
+/// 1 <= S_k <= A_k (that is, non trivial slicing occurs along exactly
+/// one dimension),
+/// S_i = A_i for all i > k (that is, all trailing dimensions are preserved
+/// in full).
+/// In other words, the slice shape S must be of the form:
+/// [ 1, 1, ..., 1, Sk, Ak + 1, Ak + 2, ...,An ]
+///
+/// In case the size and/or offset extracted are dynamic then this is possible
+/// only if there is single dimension in the reassociation group that has a size
+/// not equal to 1.
+/// In other words, the tensor shape must be of the form:
+/// [ 1, 1, ..., 1, A, 1, ...,1 ]
+/// Note - it might be possible to enable this pattern for more cases when the
+/// size/offset are dynamic via performing an analysis of the possible values
+/// that could be given to the size/offset.
+///
+/// Example:
+/// The transformation is possible because each reassociation group can be
+/// represented as a contiguous slice (i.e., [8x16->2x16], [1x7->1x?],
+/// [20->10]).
+/// ```
+/// BEFORE:
+/// %collapse = tensor.collapse_shape %src [[0, 1], [2, 3], [4]] ...
+/// tensor<8x16x1x7x20f32> to tensor<128x7x20xf32>
+/// %slice = tensor.extract_slice %slice [0, 0, 0][32, %size, 10][1, 1, 1]
+/// tensor<128x7x20xf32> to tensor<32x?x10xf32>
+///
+/// AFTER:
+/// %slice = tensor.extract_slice %src [0, 0, 0, 0, 0][2, 16, 1, %size, 10]
+// [1, 1, 1, 1, 1] : tensor<8x16x1x7x20f32> to tensor<2x16x1x?x10xf32>
+/// %collapse = tensor.collapse_shape %slice [[0, 1], [2, 3], [4]] ...
+/// tensor<2x16x1x?x10xf32> to tensor<32x?x10xf32>
+/// ```
+struct BubbleUpCollapseShapeThroughExtractSlice
+ : public OpRewritePattern<tensor::ExtractSliceOp> {
+ using OpRewritePattern<tensor::ExtractSliceOp>::OpRewritePattern;
+
+ LogicalResult matchAndRewrite(tensor::ExtractSliceOp sliceOp,
+ PatternRewriter &rewriter) const override {
+ auto collapseShapeOp =
+ sliceOp.getSource().getDefiningOp<tensor::CollapseShapeOp>();
+ if (!collapseShapeOp)
+ return rewriter.notifyMatchFailure(
+ sliceOp,
+ "tensor.extract_slice source not produced by tensor.collapse_shape");
+
+ if (!sliceOp.hasUnitStride()) {
+ return rewriter.notifyMatchFailure(
+ sliceOp, "unsupported: non-unit stride. Only contiguous slices can "
+ "be supported in this transformation.");
+ }
+
+ // The tensor.extract_slice before applying the pattern works on the result
+ // of the tensor.collapse_shape, so variables (i.e. inputs for
+ // ExtractSliceOp) referring to the state before applying the pattern are
+ // named with the prefix "collapsed", and ones referring to the state after
+ // applying the pattern are named with the prefix "expanded".
+ SmallVector<OpFoldResult> collapsedOffsets = sliceOp.getMixedOffsets();
+ SmallVector<OpFoldResult> collapsedSizes = sliceOp.getMixedSizes();
+
+ if (static_cast<size_t>(sliceOp.getResultType().getRank()) !=
+ collapsedSizes.size())
+ return rewriter.notifyMatchFailure(sliceOp,
+ "unimplemented: rank reducing slice");
+
+ ArrayRef<int64_t> srcShape = collapseShapeOp.getSrcType().getShape();
+ SmallVector<ReassociationIndices, 4> reassociationIndices =
+ collapseShapeOp.getReassociationIndices();
+
+ // Compute new offsets, sizes, and strides for tensor.extract_slice.
+ // The new tensor.extract_slice will work on a tensor that has has a rank
+ // equal to the rank of the src of the collapse_shape. In each iteration of
+ // the loop, the offsets and sizes will be computed per reassociation group.
+ SmallVector<OpFoldResult> expandedOffsets, expandedSizes;
+ SmallVector<OpFoldResult> expandedStrides(srcShape.size(),
+ rewriter.getIndexAttr(1));
+
+ for (auto [groupIdx, reassocIndices] :
+ enumerate(collapseShapeOp.getReassociationIndices())) {
+ OpFoldResult collapsedSize = collapsedSizes[groupIdx];
+ OpFoldResult collapsedOffset = collapsedOffsets[groupIdx];
+ // 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)) {
+ int nonUnitSizeCount = 0;
+ for (int64_t expandedShapeIdx : reassocIndices) {
+ if (srcShape[expandedShapeIdx] != 1) {
+ nonUnitSizeCount++;
+ expandedSizes.emplace_back(collapsedSize);
+ expandedOffsets.emplace_back(collapsedOffset);
+ continue;
+ }
+
+ expandedSizes.emplace_back(rewriter.getIndexAttr(1));
+ expandedOffsets.emplace_back(rewriter.getIndexAttr(0));
+ }
+
+ if (nonUnitSizeCount != 1) {
+ return rewriter.notifyMatchFailure(
+ sliceOp,
+ "unsupported: slice cannot be verified to be contiguous");
+ }
+ continue;
+ }
+
+ // Case #2 = size and offset are static.
+ // Verify that the slice can be represented as a contiguous slice of the
+ // src of the collapse_shape.
+ // Checking this must be done on order of most
+ // internal dimensions first, so traversal is done in reverse order of the
+ // reassociation group.
+ int64_t collapsedSizeValue = getConstantIntValue(collapsedSize).value();
+ int64_t collapsedOffsetValue =
+ getConstantIntValue(collapsedOffset).value();
+
+ SmallVector<OpFoldResult> groupExpandedSizes, groupExpandedOffsets;
+
+ for (int64_t expandedShapeIdx : llvm::reverse(reassocIndices)) {
+ int64_t expandedShapeSize = srcShape[expandedShapeIdx];
+
+ // This is a dimension that slicing will occur on, so need to make sure
+ // that the slice size can be set to the shape size and the offset to 0.
+ if (collapsedSizeValue >= expandedShapeSize &&
+ (collapsedSizeValue % expandedShapeSize != 0 ||
+ collapsedOffsetValue % expandedShapeSize != 0)) {
+ return rewriter.notifyMatchFailure(
+ sliceOp, "unsupported: cannot be extracted as a contiguous slice "
+ "of the src of the collapse_shape");
+ }
+
+ int64_t offsetInDim = collapsedOffsetValue % expandedShapeSize;
+
+ // This is the dimension that slicing will occur along, so need to make
+ // sure that the slice size + offset will not exceed the shape size.
+ if (collapsedSizeValue < expandedShapeSize &&
+ (collapsedSizeValue + offsetInDim) >= expandedShapeSize) {
+ return rewriter.notifyMatchFailure(
+ sliceOp, "unsupported: slice cannot be extracted as a contiguous "
+ "slice of the src of the collapse_shape");
+ }
+
+ groupExpandedSizes.emplace_back(rewriter.getIndexAttr(
+ std::min(collapsedSizeValue, expandedShapeSize)));
+ groupExpandedOffsets.emplace_back(rewriter.getIndexAttr(offsetInDim));
+
+ // Remove the size and offset of trailing dimensions from the size and
+ // offset of the slice.
+ collapsedSizeValue /= expandedShapeSize;
+ collapsedSizeValue = std::max<int64_t>(collapsedSizeValue, 1);
+ collapsedOffsetValue /= expandedShapeSize;
+ }
+
+ expandedSizes.append(groupExpandedSizes.rbegin(),
+ groupExpandedSizes.rend());
+ expandedOffsets.append(groupExpandedOffsets.rbegin(),
+ groupExpandedOffsets.rend());
+ }
+
+ Value newSliceOp = rewriter.create<tensor::ExtractSliceOp>(
+ collapseShapeOp->getLoc(), collapseShapeOp.getSrc(), expandedOffsets,
+ expandedSizes, expandedStrides);
+ rewriter.replaceOpWithNewOp<tensor::CollapseShapeOp>(
+ sliceOp, sliceOp.getResultType(), newSliceOp,
+ collapseShapeOp.getReassociationIndices());
+
+ return success();
+ }
+};
+
} // namespace
void mlir::tensor::populateReassociativeReshapeFoldingPatterns(
@@ -448,5 +634,6 @@ void mlir::tensor::populateBubbleUpExpandShapePatterns(
void mlir::tensor::populateBubbleUpExtractSliceOpPatterns(
RewritePatternSet &patterns) {
- patterns.add<BubbleUpExpandShapeThroughExtractSlice>(patterns.getContext());
+ patterns.add<BubbleUpExpandShapeThroughExtractSlice,
+ BubbleUpCollapseShapeThroughExtractSlice>(patterns.getContext());
}
diff --git a/mlir/test/Dialect/Linalg/transform-op-fuse.mlir b/mlir/test/Dialect/Linalg/transform-op-fuse.mlir
index 9bcc125ce1ba9..441020f1cddfc 100644
--- a/mlir/test/Dialect/Linalg/transform-op-fuse.mlir
+++ b/mlir/test/Dialect/Linalg/transform-op-fuse.mlir
@@ -438,3 +438,53 @@ module attributes {transform.with_named_sequence} {
transform.yield
}
}
+
+// -----
+
+// CHECK-LABEL: func.func @bubble_up_extract_slice_through_collapse_shape(
+// CHECK: scf.for %[[X:[A-Za-z0-9]+]] = {{.*}} -> (tensor<8x1800x32xf32>) {
+// CHECK: %[[EXTRACT1:.*]] = tensor.extract_slice
+// CHECK: %[[COLLAPSE1:.*]] = tensor.collapse_shape %[[EXTRACT1]]
+// CHECK: %[[EXP1:.*]] = linalg.exp ins(%[[COLLAPSE1]]
+func.func @bubble_up_extract_slice_through_collapse_shape(%0: tensor<1x8x1800x32xf32>) -> tensor<8x1800x32xf32> {
+ %expand = tensor.collapse_shape %0 [[0, 1], [2], [3]] : tensor<1x8x1800x32xf32> into tensor<8x1800x32xf32>
+ %empty = tensor.empty() : tensor<8x1800x32xf32>
+ %exp = linalg.exp ins(%expand : tensor<8x1800x32xf32>) outs(%empty : tensor<8x1800x32xf32>) -> tensor<8x1800x32xf32>
+ return %exp : tensor<8x1800x32xf32>
+}
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {
+ %0 = transform.structured.match ops{["linalg.exp"]} in %arg0 : (!transform.any_op) -> !transform.any_op
+ %transformed, %loops:1 = transform.structured.fuse %0 [1, 0, 0] interchange [0, 1, 2] apply_cleanup = true :
+ (!transform.any_op) -> (!transform.any_op, !transform.op<"scf.for">)
+ transform.yield
+ }
+}
+
+
+// -----
+
+// CHECK-LABEL: func.func @bubble_up_extract_slice_through_collapse_shape_with_collapse_producer(
+// CHECK: scf.for %[[X:[A-Za-z0-9]+]] = {{.*}}
+// CHECK: %[[VAL_9:.*]] = tensor.extract_slice
+// CHECK: %[[VAL_11:.*]] = linalg.abs ins(%[[VAL_9]]
+// CHECK: %[[VAL_12:.*]] = tensor.collapse_shape %[[VAL_11]]
+// CHECK: %[[VAL_14:.*]] = linalg.exp ins(%[[VAL_12]]
+func.func @bubble_up_extract_slice_through_collapse_shape_with_collapse_producer(%0: tensor<1x8x1800x32xf32>) -> tensor<8x1800x32xf32> {
+ %empty1 = tensor.empty() : tensor<1x8x1800x32xf32>
+ %abs = linalg.abs ins(%0 : tensor<1x8x1800x32xf32>) outs(%empty1 : tensor<1x8x1800x32xf32>) -> tensor<1x8x1800x32xf32>
+ %expand = tensor.collapse_shape %abs [[0, 1], [2], [3]] : tensor<1x8x1800x32xf32> into tensor<8x1800x32xf32>
+ %empty2 = tensor.empty() : tensor<8x1800x32xf32>
+ %exp = linalg.exp ins(%expand : tensor<8x1800x32xf32>) outs(%empty2 : tensor<8x1800x32xf32>) -> tensor<8x1800x32xf32>
+ return %exp : tensor<8x1800x32xf32>
+}
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {
+ %0 = transform.structured.match ops{["linalg.exp"]} in %arg0 : (!transform.any_op) -> !transform.any_op
+ %transformed, %loops:1 = transform.structured.fuse %0 [1, 0, 0] interchange [0, 1, 2] apply_cleanup = true :
+ (!transform.any_op) -> (!transform.any_op, !transform.op<"scf.for">)
+ transform.yield
+ }
+}
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 3900bc56f433d..d05bf1bf76f29 100644
--- a/mlir/test/Dialect/Tensor/bubble-up-extract-slice-op.mlir
+++ b/mlir/test/Dialect/Tensor/bubble-up-extract-slice-op.mlir
@@ -113,6 +113,159 @@ func.func @bubble_up_extract_slice_affine_apply_not_folded(%src: tensor<60xf32>,
return %extract : tensor<?x5x2xf32>
}
+// CHECK-LABEL: func.func @bubble_up_extract_slice_through_collapse_shape_single_reassoc_group(
+// CHECK-SAME: %[[SRC:.*]]: tensor<6x5x2xf32>) -> tensor<1xf32> {
+// CHECK: %[[EXTRACT:.*]] = tensor.extract_slice %[[SRC]][0, 0, 0] [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_single_reassoc_group(%src: tensor<6x5x2xf32>) -> tensor<1xf32> {
+ %collapse = tensor.collapse_shape %src [[0, 1, 2]] : tensor<6x5x2xf32> into tensor<60xf32>
+ %extract = tensor.extract_slice %collapse[0][1][1] : tensor<60xf32> to tensor<1xf32>
+ return %extract : tensor<1xf32>
+}
+
+// CHECK-LABEL: func.func @bubble_up_extract_slice_through_collapse_shape_multiple_reassoc_group(
+// CHECK-SAME: %[[VAL_0:.*]]: tensor<6x5x3x10xf32>) -> tensor<15x10xf32> {
+// CHECK: %[[VAL_1:.*]] = tensor.extract_slice %[[VAL_0]][1, 0, 1, 0] [3, 5, 1, 10] [1, 1, 1, 1]
+// CHECK: %[[VAL_2:.*]] = tensor.collapse_shape %[[VAL_1]] {{\[\[}}0, 1], [2, 3]]
+// CHECK: return %[[VAL_2]]
+func.func @bubble_up_extract_slice_through_collapse_shape_multiple_reassoc_group(%src: tensor<6x5x3x10xf32>) -> tensor<15x10xf32> {
+ %collapse = tensor.collapse_shape %src [[0, 1], [2, 3]] : tensor<6x5x3x10xf32> into tensor<30x30xf32>
+ %extract = tensor.extract_slice %collapse[5, 10][15, 10][1, 1] : tensor<30x30xf32> to tensor<15x10xf32>
+ return %extract : tensor<15x10xf32>
+}
+
+// CHECK-LABEL: func.func @bubble_up_extract_slice_through_collapse_shape_offset_on_leading_dim(
+// CHECK-SAME: %[[VAL_0:.*]]: tensor<6x5x2xf32>) -> tensor<4xf32> {
+// CHECK: %[[VAL_1:.*]] = tensor.extract_slice %[[VAL_0]][2, 0, 0] [1, 2, 2] [1, 1, 1]
+// CHECK: %[[VAL_2:.*]] = tensor.collapse_shape %[[VAL_1]] {{\[\[}}0, 1, 2]]
+// CHECK: return %[[VAL_2]]
+func.func @bubble_up_extract_slice_through_collapse_shape_offset_on_leading_dim(%src: tensor<6x5x2xf32>) -> tensor<4xf32> {
+ %collapse = tensor.collapse_shape %src [[0, 1, 2]] : tensor<6x5x2xf32> into tensor<60xf32>
+ %extract = tensor.extract_slice %collapse[20][4][1] : tensor<60xf32> to tensor<4xf32>
+ return %extract : tensor<4xf32>
+}
+
+// CHECK-LABEL: func.func @bubble_up_extract_slice_through_collapse_shape_dynamic_size(
+// CHECK-SAME: %[[SRC:.*]]: tensor<1x5x1xf32>,
+// CHECK-SAME: %[[SIZE:.*]]: index) -> tensor<?xf32> {
+// CHECK: %[[EXTRACT:.*]] = tensor.extract_slice %[[SRC]][0, 0, 0] [1, %[[SIZE]], 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_dynamic_size(%src: tensor<1x5x1xf32>, %size : index) -> tensor<?xf32> {
+ %collapse = tensor.collapse_shape %src [[0, 1, 2]] : tensor<1x5x1xf32> into tensor<5xf32>
+ %extract = tensor.extract_slice %collapse[0][%size][1] : tensor<5xf32> to tensor<?xf32>
+ return %extract : tensor<?xf32>
+}
+
+// CHECK-LABEL: func.func @bubble_up_extract_slice_through_collapse_shape_dynamic_size_and_src(
+// CHECK-SAME: %[[SRC:.*]]: tensor<1x?x1xf32>,
+// CHECK-SAME: %[[SIZE:.*]]: index) -> tensor<?xf32> {
+// CHECK: %[[EXTRACT:.*]] = tensor.extract_slice %[[SRC]][0, 0, 0] [1, %[[SIZE]], 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_dynamic_size_and_src(%src: tensor<1x?x1xf32>, %size : index) -> tensor<?xf32> {
+ %collapse = tensor.collapse_shape %src [[0, 1, 2]] : tensor<1x?x1xf32> into tensor<?xf32>
+ %extract = tensor.extract_slice %collapse[0][%size][1] : tensor<?xf32> to tensor<?xf32>
+ 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> {
+// CHECK: %[[EXTRACT:.*]] = tensor.extract_slice %[[SRC]][0, %[[OFFSET]], 0] [1, 3, 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_dynamic_offset(%src: tensor<1x5x1xf32>, %offset : index) -> tensor<3xf32> {
+ %collapse = tensor.collapse_shape %src [[0, 1, 2]] : tensor<1x5x1xf32> into tensor<5xf32>
+ %extract = tensor.extract_slice %collapse[%offset][3][1] : tensor<5xf32> to tensor<3xf32>
+ return %extract : tensor<3xf32>
+}
+
+// CHECK-LABEL: func.func @bubble_up_extract_slice_through_collapse_shape_dynamic_offset_and_size(
+// CHECK-SAME: %[[SRC:.*]]: tensor<14x1xf32>,
+// CHECK-SAME: %[[OFFSET:.*]]: index,
+// CHECK-SAME: %[[SIZE:.*]]: index) -> tensor<?xf32> {
+// CHECK: %[[EXTRACT:.*]] = tensor.extract_slice %[[SRC]]{{\[}}%[[OFFSET]], 0] {{\[}}%[[SIZE]], 1] [1, 1]
+// CHECK: %[[COLLAPSE:.*]] = tensor.collapse_shape %[[EXTRACT]] {{\[\[}}0, 1]]
+// CHECK: return %[[COLLAPSE]]
+func.func @bubble_up_extract_slice_through_collapse_shape_dynamic_offset_and_size(%src: tensor<14x1xf32>, %offset : index, %size : index) -> tensor<?xf32> {
+ %collapse = tensor.collapse_shape %src [[0, 1]] : tensor<14x1xf32> into tensor<14xf32>
+ %extract = tensor.extract_slice %collapse[%offset][%size][1] : tensor<14xf32> to tensor<?xf32>
+ return %extract : tensor<?xf32>
+}
+
+// CHECK-LABEL: func.func @bubble_up_extract_slice_through_collapse_shape_dynamic_and_static_groups(
+// CHECK-SAME: %[[SRC:.*]]: tensor<5x10x1x1x40xf32>,
+// CHECK-SAME: %[[OFFSET:.*]]: index,
+// CHECK-SAME: %[[SIZE:.*]]: index) -> tensor<20x?xf32> {
+// CHECK: %[[EXTRACT:.*]] = tensor.extract_slice %[[SRC]][1, 0, 0, 0, %[[OFFSET]]] [2, 10, 1, 1, %[[SIZE]]] [1, 1, 1, 1, 1]
+// CHECK: %[[COLLAPSE:.*]] = tensor.collapse_shape %[[EXTRACT]] {{\[\[}}0, 1], [2, 3, 4]]
+// CHECK: return %[[COLLAPSE]]
+func.func @bubble_up_extract_slice_through_collapse_shape_dynamic_and_static_groups(%src: tensor<5x10x1x1x40xf32>, %offset : index, %size : index) -> tensor<20x?xf32> {
+ %collapse = tensor.collapse_shape %src [[0, 1], [2, 3, 4]] : tensor<5x10x1x1x40xf32> into tensor<50x40xf32>
+ %extract = tensor.extract_slice %collapse[10, %offset][20...
[truncated]
|
|
@qedawkins @banach-space I've added a pattern similar to the one from PR #126898 but this time for tensor.collpase_shape. |
Thanks! I’ve been a bit distracted, but I’ll definitely review this before the weekend. I see a lot of comments already, so I’m liking this 🙂 By the way, is this being upstreamed from IREE, or was it implemented from scratch? |
Thanks for reviewing! |
Comment on lines
+439
to
+440
| /// In case the size and offset extracted are static then this is possible if | ||
| /// the following conditions are met within each reassociation group: |
There was a problem hiding this comment.
I struggle with this part a bit. Shouldn't there be only one reassociation group in which slicing happens? And, within that group, exactly one expanded dim should be sliced? If I am incorrect, is there an example?
There was a problem hiding this comment.
I'll start out with the second part of your question regarding each individual reassociation group.
We need to make sure that after we bubble up the tensor.extract_slice that the data we extract from the expanded tensor for this group is contiguous, since it was contiguous before the bubble up.
Now lets look at the post bubble up sizes [ 1, 1, ..., 1, Sk, Ak + 1, Ak + 2, ...,An ] and see why this makes sense.
If we start with the trailing dims, we want them all the be full, which means that the data is obviously contiguous. Then we have this one dim where we extract a slice that is smaller than the full size Sk, but things are still contiguous since we now have a contiguous slice of size Sk*(Prod(Ai) for i=k+1 to n). Note this this slice might not be with offset 0, but that is still fine since the original slice we extracted might not have been with offset 0.
Then we have all those leading dime of size 1. Since they are size 1 this means that we could say that we are technically slicing on this dim but in practice the size 1 means that we aren't breaking contiguity and just maybe changing the offset of the contiguous block that we extract.
to summarize - we could technically say that we are slicing on k different dims, but in practice the result is a single contiguous block of data as required.
Now regarding the first part of the question about different reassociation groups.
The data extracted from the collapsed tensor might have been sliced on multiple different collapsed dims which could make the entire original slice to not be contiguous, so it should be fine if slicing occurs on multiple different reassociation groups after the bubble up.
I'll give an example for this from the tests:
BEFORE:
%collapse = tensor.collapse_shape %src [[0, 1], [2, 3]] : tensor<6x5x3x10xf32> into tensor<30x30xf32>
%extract = tensor.extract_slice %collapse[5, 10][15, 10][1, 1] : tensor<30x30xf32> to tensor<15x10xf32>
AFTER:
%extracted_slice = tensor.extract_slice %arg0[1, 0, 1, 0] [3, 5, 1, 10] [1, 1, 1, 1] : tensor<6x5x3x10xf32> to tensor<3x5x1x10xf32>
%collapsed = tensor.collapse_shape %extracted_slice [[0, 1], [2, 3]] : tensor<3x5x1x10xf32> into tensor<15x10xf32>
In the example there are 2 different reassociation groups, where a slice is extracted from each one of them, and the result of the bubble up is still legal and the result of the tensor.collapse_shape post bubble up represents the same data as the result of the tensor.extract_slice pre bubble up.
We can also directly calculate the indices within the src tensor of the elements from the result and see that they are equal in both cases: [160:169], [190:199], [220:229]...
There was a problem hiding this comment.
Thank you for taking the time to write this very comprehensive explanation, this makes it super clear to me 🙏🏻 (and I was indeed incorrect)
I was just about to ask you to add this example as a test, but that's already done! :)
| /// AFTER: | ||
| /// %slice = tensor.extract_slice %src [0, 0, 0, 0, 0][2, 16, 1, %size, 10] | ||
| // [1, 1, 1, 1, 1] : tensor<8x16x1x7x20f32> to tensor<2x16x1x?x10xf32> | ||
| /// %collapse = tensor.collapse_shape %slice [[0, 1], [2, 3], [4]] ... |
There was a problem hiding this comment.
Whew! This is actually correct. Took me a while to work out that it, but this works because for the reassociation the other dimensions are all 1.
banach-space
approved these changes
Mar 25, 2025
There was a problem hiding this comment.
I've left a few minor comments/nits — nothing blocking. Overall, this is a very high-quality submission — thank you for all the effort! LGTM.
@MaheshRavishankar, this falls much closer to your area of expertise, so thanks for taking a look!
Comment on lines
+439
to
+440
| /// In case the size and offset extracted are static then this is possible if | ||
| /// the following conditions are met within each reassociation group: |
There was a problem hiding this comment.
Thank you for taking the time to write this very comprehensive explanation, this makes it super clear to me 🙏🏻 (and I was indeed incorrect)
I was just about to ask you to add this example as a test, but that's already done! :)
Thank you very much for taking the time to review this! |
|
@MaheshRavishankar , is this ready to land? |
|
I'll take a look at this tomorrow cause I need to understand this a bit better |
Hi @MaheshRavishankar have you had a chance to review? |
MaheshRavishankar
approved these changes
Apr 2, 2025
|
|
||
| if (static_cast<size_t>(sliceOp.getResultType().getRank()) != | ||
| collapsedSizes.size()) | ||
| return rewriter.notifyMatchFailure(sliceOp, |
There was a problem hiding this comment.
Nit: Please add { } around multi-line statements.
There was a problem hiding this comment.
Added in all the relevant places in the pattern
| SmallVector<OpFoldResult> expandedStrides(srcShape.size(), | ||
| rewriter.getIndexAttr(1)); | ||
|
|
||
| for (auto [groupIdx, reassocIndices] : |
There was a problem hiding this comment.
Nit: You could do
for (auto [collapsedSize, collapsedOffset, reassocIndicess] : llvm::zip_equal(collapsedSizes, collapsedOffsets, collapsedShapeOp.getReassociationIndices())
2f05786 to
3c69390
Compare
Thanks for the review! |
|
LLVM Buildbot has detected a new failure on builder Full details are available at: https://lab.llvm.org/buildbot/#/builders/55/builds/9329 Here is the relevant piece of the build log for the reference |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Add a pattern that bubbles up tensor.extract_slice through tensor.collapse_shape.
The pattern is registered in a pattern population function that is used by the transform op
transform.apply_patterns.tensor.bubble_up_extract_slice and by the tranform op transform.structured.fuse as a cleanup pattern.
This pattern enables tiling and fusing op chains which contain tensor.collapse_shape if added as a cleanup pattern of tile and fuse utility.
Without this pattern that would not be possible, as tensor.collapse_shape does not implement the tiling interface. This is an additional pattern to the one added in PR #126898