[PATCH] D142557: [mlir][Linalg] Fix crash when converting rank-reduced tensor.insert_slice

[Cullen Rhodes via Phabricator via llvm-commits]

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>