[mlir][tensor] Keep size-1 dim for non-rank-reducing extract_slice

[antiagainst]

Summary

This commit fixes a crash for cases like:

%0 = tensor.extract_slice %arg0[0, 0] [%c1, %size] [1, 1] :
  tensor<?x?xf32> to tensor<?x?xf32>

Previous the first size-1 dimension will be dropped, thus causing
an mismatch in tensor rank.

Depends On D117020

[mravishankar]

This by itself is fine, but it might be hiding an error on the caller side. Can you tell me where the call site is when the crash happens.

mlir/lib/Dialect/Tensor/IR/TensorOps.cpp

getCanonicalSliceResultType(unsigned resultRank, RankedTensorType sourceType,
   }
   return resultType;
 }
 
 llvm::SmallDenseSet<unsigned> ExtractSliceOp::getDroppedDims() {
   llvm::SmallDenseSet<unsigned> droppedDims;
   ArrayRef<int64_t> resultShape = getType().getShape();
   SmallVector<OpFoldResult> mixedSizes = getMixedSizes();
+  if (resultShape.size() == mixedSizes.size())
+    return {};
   unsigned shapePos = 0;
   for (const auto &size : enumerate(mixedSizes)) {
     Optional<int64_t> sizeVal = getConstantIntValue(size.value());
     // If the size is not 1, or if the current matched dimension of the result
     // is the same static shape as the size value (which is 1), then the
     // dimension is preserved.
     if (!sizeVal || sizeVal.getValue() != 1 ||
         (shapePos < resultShape.size() && resultShape[shapePos] == 1)) {

mlir/test/Dialect/Tensor/resolve-shaped-type-result-dims.mlir

func @pad_both_low_and_high_pad(%tensor: tensor<1x56x56x144xf32>, %arg0: index, %arg1: index, %arg2: index) -> (index, index) {
   %dim1 = tensor.dim %pad, %c1 : tensor<1x?x?x16xf32>
   %dim2 = tensor.dim %pad, %c2 : tensor<1x?x?x16xf32>
   return %dim1, %dim2 : index, index
 }
 
 // CHECK-LABEL: func @pad_both_low_and_high_pad
 //       CHECK:   %[[C4:.+]] = arith.constant 4 : index
 //       CHECK:   return %[[C4]], %[[C4]]
+
+// -----
+
+func @extract_slice_same_rank(%arg0 : tensor<?x?xf32>,
+    %arg1 : index, %arg2: index) -> (index, index) {
+  %f0 = arith.constant 0.0 : f32
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %0 = tensor.extract_slice %arg0[0, 0] [%c1, %arg1] [1, 1] :
+     tensor<?x?xf32> to tensor<?x?xf32>
+  %pad = tensor.pad %0 low[0, 0] high[0, %arg2] {
+  ^bb0(%arg3: index, %arg4: index):
+    tensor.yield %f0 : f32
+  } : tensor<?x?xf32> to tensor<?x?xf32>
+  %1 = tensor.dim %pad, %c0 : tensor<?x?xf32>
+  %2 = tensor.dim %pad, %c1 : tensor<?x?xf32>
+  return %1, %2 : index, index
+}
+
+//      CHECK: #[[MAP:.+]] = affine_map<()[s0, s1] -> (s1 + s0)>
+//      CHECK: func @extract_slice_same_rank
+// CHECK-SAME: %{{.+}}: tensor<?x?xf32>, %[[ARG1:.+]]: index, %[[ARG2:.+]]: index
+//      CHECK:   %[[DIM0:.+]] = arith.constant 1 : index
+//      CHECK:   %[[DIM1:.+]] = affine.apply #[[MAP]]()[%[[ARG2]], %[[ARG1]]]
+//      CHECK:   return %[[DIM0]], %[[DIM1]]