[mlir][tensor] Fix bug in utility `tensor::isCastLikeExtractSliceOp`

[christopherbate]

Summary

Fixes an issue where isCastLikeExtractSliceOp did not account for the fact
that tensor.extract_slice may drop non-unit dimensions. This change makes the
utility function behave inline with its name/description. The only user of this
function is in the FindPayloadReplacementOpInterface for the
tensor::ExtractSliceOp. This can potentially cause downstream projects to have
more "listener could not find replacement op" errors when interpreting Transform
IR, but the behavior is inline with the documented conservative behavior of the
Transform dialect's TrackingListener.

[springerm]

mlir/lib/Dialect/Tensor/Utils/Utils.cpp

96

no need to change this here, but I just noticed that CastLike is not really the correct term. It should be CastOrReshapeLike (or something like that...)

mlir/include/mlir/Dialect/Tensor/Utils/Utils.h

computeTransposedType(RankedTensorType rankedTensorType,
                       ArrayRef<int64_t> transposeVector);
 
 /// A tensor.insert_slice is a cast-like operation if it merely rank-extends the
 /// source tensor or inserts the source tensor into a destination tensor with
 /// the same shape.
 bool isCastLikeInsertSliceOp(InsertSliceOp op);
 
 /// A tensor.extract_slice is a cast-like operation if it merely rank-reduces
-/// the source tensor or extracts the entire source tensor.
+/// unit dimensions of the source tensor or extracts the entire source tensor.
 bool isCastLikeExtractSliceOp(ExtractSliceOp op);
 
 } // namespace tensor
 } // namespace mlir
 
 #endif // MLIR_DIALECT_TENSOR_UTILS_UTILS_H_

mlir/lib/Dialect/Tensor/Utils/Utils.cpp

for (int64_t resultDim = 0; resultDim < op.getDestType().getRank();
     if (failed(equalDimSize) || !*equalDimSize)
       return false;
     ++srcDim;
   }
 
   return true;
 }
 
 bool mlir::tensor::isCastLikeExtractSliceOp(ExtractSliceOp op) {
   llvm::SmallBitVector droppedDims = op.getDroppedDims();
   int64_t resultDim = 0;
   // Source dims and result dims (apart from dropped dims) must have the same
   // size.
-  for (int64_t dim = 0; dim < op.getSourceType().getRank(); ++dim) {
+  RankedTensorType sourceType = op.getSourceType();
+  for (int64_t dim = 0, e = sourceType.getRank(); dim < e; ++dim) {
     if (droppedDims.test(dim)) {
+      // ExtractSlice may drop unit dimensions that result from taking a size-1
+      // slice from a non-size-1 source dimension.
+      if (sourceType.getDimSize(dim) != 1)
+        return false;
       continue;
     }
     FailureOr<bool> equalDimSize = ValueBoundsConstraintSet::areEqual(
         op.getSource(), op.getResult(), dim, resultDim);
     if (failed(equalDimSize) || !*equalDimSize)
       return false;
     ++resultDim;
   }
 
   return true;
 }

mlir/test/Dialect/Tensor/tracking-listener.mlir

 func.func @non_cast_like_extract_slice() {
   // expected-error @below {{listener could not find replacement op}}
   %0 = "test.foo"() {replaced} : () -> (tensor<5xf32>)
   %1 = "test.foo"() : () -> (tensor<1x5x1x1xf32>)
   %2 = tensor.extract_slice %1[0, 0, 0, 0][1, 3, 1, 1][1, 1, 1, 1]
       {replacement_0 = 0} : tensor<1x5x1x1xf32> to tensor<3xf32>
   return
 }
+
+// -----
+
+func.func @non_cast_like_extract_slice_drop_non_unit_dim() {
+  // expected-error @below {{listener could not find replacement op}}
+  %0 = "test.foo"() {replaced} : () -> (tensor<f32>)
+  %1 = "test.foo"() : () -> (tensor<1x5x1x1xf32>)
+  %2 = tensor.extract_slice %1[0, 0, 0, 0][1, 1, 1, 1][1, 1, 1, 1]
+      {replacement_0 = 0} : tensor<1x5x1x1xf32> to tensor<f32>
+  return
+}