[PATCH] D142557: [mlir][Linalg] Fix crash when converting rank-reduced tensor.insert_slice
Cullen Rhodes via Phabricator via llvm-commits
llvm-commits at lists.llvm.org
Wed Jan 25 08:36:02 PST 2023
c-rhodes created this revision.
c-rhodes added reviewers: mravishankar, nicolasvasilache.
Herald added subscribers: hanchung, Moerafaat, bzcheeseman, sdasgup3, wenzhicui, wrengr, jsetoain, cota, teijeong, rdzhabarov, tatianashp, msifontes, jurahul, Kayjukh, grosul1, Joonsoo, liufengdb, aartbik, mgester, arpith-jacob, antiagainst, shauheen, rriddle, mehdi_amini.
Herald added a reviewer: hanchung.
Herald added a project: All.
c-rhodes requested review of this revision.
Herald added subscribers: limo1996, stephenneuendorffer.
Herald added a project: MLIR.
In conversion of tensor.insert_slice op with leading dimensions of size
1 to rank-reduced version, an out-of-bounds call to ArrayRef::slice
causes a crash. It looks like the reassociation is incorrectly computed
for tensors with leading dimensions of size 1.
I'm sure there's a better way to fix this, input would be welcome.
See https://github.com/llvm/llvm-project/issues/60291 for more info.
Repository:
rG LLVM Github Monorepo
https://reviews.llvm.org/D142557
Files:
mlir/include/mlir/Dialect/Linalg/Utils/Utils.h
mlir/lib/Dialect/Linalg/Transforms/DropUnitDims.cpp
mlir/lib/Dialect/Linalg/Utils/Utils.cpp
mlir/test/Dialect/Linalg/drop-unit-extent-dims.mlir
Index: mlir/test/Dialect/Linalg/drop-unit-extent-dims.mlir
===================================================================
--- mlir/test/Dialect/Linalg/drop-unit-extent-dims.mlir
+++ mlir/test/Dialect/Linalg/drop-unit-extent-dims.mlir
@@ -501,6 +501,18 @@
// -----
+func.func @insert_slice_lower_rank(%arg0: tensor<1x1x32xi32>, %arg1: tensor<1x32xi32>) -> tensor<1x1x32xi32> {
+ %0 = tensor.insert_slice %arg1 into %arg0[0, 0, 0] [1, 1, 32] [1, 1, 1] : tensor<1x32xi32> into tensor<1x1x32xi32>
+ return %0 : tensor<1x1x32xi32>
+}
+// CHECK-LABEL: func @insert_slice_lower_rank
+// CHECK: %[[RESHAPE:.+]] = tensor.collapse_shape %{{.+}} {{\[}}[0, 1]] : tensor<1x32xi32> into tensor<32xi32>
+// CHECK: %[[RESULT:.+]] = tensor.insert_slice %[[RESHAPE]]
+// CHECK-SAME: tensor<32xi32> into tensor<1x1x32xi32>
+// CHECK: return %[[RESULT]]
+
+// -----
+
#accesses = [
affine_map<(i, j, k, l, m) -> (i, k, m)>,
affine_map<(i, j, k, l, m) -> ()>,
Index: mlir/lib/Dialect/Linalg/Utils/Utils.cpp
===================================================================
--- mlir/lib/Dialect/Linalg/Utils/Utils.cpp
+++ mlir/lib/Dialect/Linalg/Utils/Utils.cpp
@@ -1039,16 +1039,21 @@
/// but non-zero offsets are not handled by SPIR-V backend at this point (and
/// potentially cannot be handled).
std::optional<SmallVector<ReassociationIndices>>
-getReassociationMapForFoldingUnitDims(ArrayRef<OpFoldResult> mixedSizes) {
+getReassociationMapForFoldingUnitDims(ArrayRef<OpFoldResult> mixedSizes,
+ bool insert) {
SmallVector<ReassociationIndices> reassociation;
ReassociationIndices curr;
for (const auto &it : llvm::enumerate(mixedSizes)) {
auto dim = it.index();
auto size = it.value();
- curr.push_back(dim);
+ if (!insert)
+ curr.push_back(dim);
auto attr = size.dyn_cast<Attribute>();
- if (attr && attr.cast<IntegerAttr>().getInt() == 1)
+ if (attr && attr.cast<IntegerAttr>().getInt() == 1) {
+ if (insert)
+ curr.push_back(dim);
continue;
+ }
reassociation.emplace_back(ReassociationIndices{});
std::swap(reassociation.back(), curr);
}
Index: mlir/lib/Dialect/Linalg/Transforms/DropUnitDims.cpp
===================================================================
--- mlir/lib/Dialect/Linalg/Transforms/DropUnitDims.cpp
+++ mlir/lib/Dialect/Linalg/Transforms/DropUnitDims.cpp
@@ -634,7 +634,8 @@
SmallVector<OpFoldResult> offsets = insertSliceOp.getMixedOffsets();
SmallVector<OpFoldResult> sizes = insertSliceOp.getMixedSizes();
SmallVector<OpFoldResult> strides = insertSliceOp.getMixedStrides();
- auto reassociation = getReassociationMapForFoldingUnitDims(sizes);
+ auto reassociation =
+ getReassociationMapForFoldingUnitDims(sizes, /*insert=*/true);
if (!reassociation ||
reassociation->size() == static_cast<size_t>(sourceType.getRank()))
return failure();
Index: mlir/include/mlir/Dialect/Linalg/Utils/Utils.h
===================================================================
--- mlir/include/mlir/Dialect/Linalg/Utils/Utils.h
+++ mlir/include/mlir/Dialect/Linalg/Utils/Utils.h
@@ -136,7 +136,8 @@
/// but non-zero offsets are not handled by SPIR-V backend at this point (and
/// potentially cannot be handled).
std::optional<SmallVector<ReassociationIndices>>
-getReassociationMapForFoldingUnitDims(ArrayRef<OpFoldResult> mixedSizes);
+getReassociationMapForFoldingUnitDims(ArrayRef<OpFoldResult> mixedSizes,
+ bool insert = false);
/// Return the identity numeric value associated to the give op. Return
/// std::nullopt if there is no known neutral element.
-------------- next part --------------
A non-text attachment was scrubbed...
Name: D142557.492133.patch
Type: text/x-patch
Size: 3718 bytes
Desc: not available
URL: <http://lists.llvm.org/pipermail/llvm-commits/attachments/20230125/a4166dad/attachment.bin>
More information about the llvm-commits
mailing list