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

[c-rhodes]

Summary

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.

[mravishankar]

mlir/include/mlir/Dialect/Linalg/Utils/Utils.h

139

Could you document what the insert argument is doing. I am not following whats happening here.

[mravishankar]

I think the core logic being fixed here is not something I worked on, so I dont have full context. Adding folks who worked on this as reviewers...

[pifon2a]

In D142557#4083304, @mravishankar wrote:

I think the core logic being fixed here is not something I worked on, so I dont have full context. Adding folks who worked on this as reviewers...

I did not work on this.

mlir/include/mlir/Dialect/Linalg/Utils/Utils.h

 
 /// Get the reassociation maps to fold the result of a extract_slice (or source
 /// of a insert_slice) operation with given offsets, and sizes to its
 /// rank-reduced version. This is only done for the cases where the size is 1
 /// and offset is 0. Strictly speaking the offset 0 is not required in general,
 /// 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.
 std::optional<Attribute> getNeutralElement(Operation *op);
 
 //===----------------------------------------------------------------------===//
 // Fusion / Tiling utilities
 //===----------------------------------------------------------------------===//

mlir/lib/Dialect/Linalg/Transforms/DropUnitDims.cpp

struct RankReducedInsertSliceOp : public OpRewritePattern<InsertOpTy> {
   using OpRewritePattern<InsertOpTy>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(InsertOpTy insertSliceOp,
                                 PatternRewriter &rewriter) const override {
     RankedTensorType sourceType = insertSliceOp.getSourceType();
     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();
     Location loc = insertSliceOp.getLoc();
     tensor::CollapseShapeOp reshapedSource;
     {
       OpBuilder::InsertionGuard g(rewriter);
       // The only difference between InsertSliceOp and ParallelInsertSliceOp is

mlir/lib/Dialect/Linalg/Utils/Utils.cpp

 
 /// Get the reassociation maps to fold the result of a extract_slice (or source
 /// of a insert_slice) operation with given offsets, and sizes to its
 /// rank-reduced version. This is only done for the cases where the size is 1
 /// and offset is 0. Strictly speaking the offset 0 is not required in general,
 /// 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();
+    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);
   }
   // When the reassociations are not empty, then fold the remaining
   // unit-dimensions into the last dimension.  If the reassociations so far is
   // empty, then leave it emtpy. This will fold everything to a rank-0 tensor.
   if (!curr.empty() && !reassociation.empty())
     reassociation.back().append(curr.begin(), curr.end());

mlir/test/Dialect/Linalg/drop-unit-extent-dims.mlir

 // CHECK-LABEL: func @insert_slice_unit_dims
 //       CHECK:   %[[RESHAPE:.+]] = tensor.collapse_shape %{{.+}} []
 //       CHECK:   %[[RESULT:.+]] = tensor.insert_slice %[[RESHAPE]]
 //  CHECK-SAME:     tensor<f32> into tensor<1x3xf32>
 //       CHECK:   return %[[RESULT]]
 
 // -----
 
+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) -> ()>,
   affine_map<(i, j, k, l, m) -> (i, k, j, l, m)>
 ]
 
 #trait = {
   iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"],