[mlir][tensor] TrackingListener: Find replacement ops through cast-like ExtractSliceOps

[springerm]

Summary

Certain ExtractSliceOps, that do extract all elements from the destination, are treated like casts when looking for replacement ops. Such ExtractSliceOps are typically rank expansions.

Depends On: D151800

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.
+bool isCastLikeExtractSliceOp(ExtractSliceOp op);
+
 } // namespace tensor
 } // namespace mlir
 
 #endif // MLIR_DIALECT_TENSOR_UTILS_UTILS_H_

mlir/lib/Dialect/Tensor/TransformOps/TensorTransformOps.cpp

if (Operation *replacement =
             transform::TrackingListener::findReplacementOp(op, values))
       return replacement;
 
     Operation *defOp = getCommonDefiningOp(values);
     if (!defOp)
       return nullptr;
 
     // Skip cast-like operations.
-    // TODO: CastOpInterface could be used if CollapseShapeOp and ExpandShapeOp
-    // implement that interface
     values.clear();
     llvm::TypeSwitch<Operation *>(defOp)
         .Case<CastOp>([&](CastOp op) { values.push_back(op.getSource()); })
         .Case<CollapseShapeOp>(
             [&](CollapseShapeOp op) { values.push_back(op.getSrc()); })
         .Case<ExpandShapeOp>(
             [&](ExpandShapeOp op) { values.push_back(op.getSrc()); })
         .Case<ReshapeOp>(
             [&](ReshapeOp op) { values.push_back(op.getSource()); })
         .Case<InsertSliceOp>([&](InsertSliceOp op) {
           if (isCastLikeInsertSliceOp(op))
             values.push_back(op.getSource());
         })
+        .Case<ExtractSliceOp>([&](ExtractSliceOp op) {
+          if (isCastLikeExtractSliceOp(op))
+            values.push_back(op.getSource());
+        })
         .Default([](Operation *op) {});
   } while (!values.empty());
 
   return nullptr;
 }
 
 //===----------------------------------------------------------------------===//
 // MakeLoopIndependentOp

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

FailureOr<bool> equalDimSize = ValueBoundsConstraintSet::areEqual(
         op.getSource(), op.getResult(), srcDim, resultDim);
     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) {
+    if (droppedDims.test(dim)) {
+      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 @cast_like_insert_slice_dynamic(
   // Rank expansion
   // Throw in a wrench: Do not use %dim1 directly, but another SSA value that
   // has the same runtime value.
   %dim1b = tensor.dim %1, %c0 : tensor<?xf32>
   %3 = tensor.insert_slice %2 into %t[0, 0, 0][1, %dim1b, 1][1, 1, 1]
       {replacement_0 = 0} : tensor<?xf32> into tensor<1x?x1xf32>
   return
 }
+
+// -----
+
+func.func @cast_like_extract_slice() {
+  %0 = "test.foo"() {replaced} : () -> (tensor<5xf32>)
+  // expected-remark @below {{replacement found}}
+  %1 = "test.foo"() : () -> (tensor<1x5x1x1xf32>)
+  %2 = tensor.extract_slice %1[0, 0, 0, 0][1, 5, 1, 1][1, 1, 1, 1]
+      {replacement_0 = 0} : tensor<1x5x1x1xf32> to tensor<5xf32>
+  return
+}
+
+// -----
+
+func.func @cast_like_extract_slice_dynamic() {
+  %0 = "test.foo"() {replaced} : () -> (tensor<?xf32>)
+  // expected-remark @below {{replacement found}}
+  %1 = "test.foo"() : () -> (tensor<1x?x1x1xf32>)
+  %c1 = arith.constant 1 : index
+  %dim = tensor.dim %1, %c1 : tensor<1x?x1x1xf32>
+  %2 = tensor.extract_slice %1[0, 0, 0, 0][1, %dim, 1, 1][1, 1, 1, 1]
+      {replacement_0 = 0} : tensor<1x?x1x1xf32> to tensor<?xf32>
+  return
+}
+
+// -----
+
+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
+}