[mlir][tensor] TrackingListener: Support cast-like InsertSliceOps with dynamic shape

[springerm]

Summary

When looking for payload op replacements, rank-expanding InsertSliceOps of dynamically-typed tensors are now supported.

Depends On: D151443

[nicolasvasilache]

Nice!

mlir/lib/Dialect/Tensor/TransformOps/CMakeLists.txt

 add_mlir_dialect_library(MLIRTensorTransformOps
   TensorTransformOps.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Tensor/TransformOps
 
   DEPENDS
   MLIRTensorTransformOpsIncGen
 
   LINK_LIBS PUBLIC
   MLIRAffineDialect
   MLIRIR
   MLIRSCFDialect
   MLIRTensorDialect
   MLIRTensorTransforms
   MLIRTransformDialect
+  MLIRValueBoundsOpInterface
 )

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

 //===- TensorTransformOps.cpp - Implementation of tensor transform ops ----===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Tensor/TransformOps/TensorTransformOps.h"
 
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Dialect/Tensor/Transforms/Transforms.h"
 #include "mlir/Dialect/Transform/IR/TransformDialect.h"
 #include "mlir/Dialect/Transform/IR/TransformInterfaces.h"
+#include "mlir/Interfaces/ValueBoundsOpInterface.h"
 #include "llvm/ADT/TypeSwitch.h"
 
 using namespace mlir;
 using namespace tensor;
 
 //===----------------------------------------------------------------------===//
 // TrackingListener
 //===----------------------------------------------------------------------===//
 
-/// A tensor.insert_slice is a cast-like operation if it the source tensor and
-/// the destination tensor have the same number of elements. I.e., the result
-/// tensor data equals the source tensor data, maybe rank-extended to a
-/// different shape.
+/// 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.
 static bool isCastLikeInsertSliceOp(InsertSliceOp op) {
-  // TODO: Support dynamically shaped tensors. Utilize ValueBoundsOpInterface
-  // to check if source and destination have the same shape.
-  if (!op.getSourceType().hasStaticShape() ||
-      !op.getDestType().hasStaticShape())
+  llvm::SmallBitVector droppedDims = op.getDroppedDims();
+  int64_t srcDim = 0;
+  // Source dims and destination dims (apart from dropped dims) must have the
+  // same size.
+  for (int64_t resultDim = 0; resultDim < op.getDestType().getRank();
+       ++resultDim) {
+    if (droppedDims.test(resultDim)) {
+      continue;
+    }
+    FailureOr<bool> equalDimSize = ValueBoundsConstraintSet::areEqual(
+        op.getSource(), op.getResult(), srcDim, resultDim);
+    if (failed(equalDimSize) || !*equalDimSize)
       return false;
-  return op.getSourceType().getNumElements() ==
-         op.getDestType().getNumElements();
+    ++srcDim;
+  }
+
+  return true;
 }
 
 Operation *
 tensor::TrackingListener::findReplacementOp(Operation *op,
                                             ValueRange newValues) const {
   SmallVector<Value> values(newValues.begin(), newValues.end());
   do {
     if (Operation *replacement =

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

func.func @non_cast_like_insert_slice(%t: tensor<7xf32>) {
   %0 = "test.foo"() {replaced} : () -> (tensor<5xf32>)
   %1 = "test.foo"() : () -> (tensor<5xf32>)
   // This is not a cast-like insert_slice op because elements from %t are
   // contained in %2.
   %2 = tensor.insert_slice %1 into %t[0][5][1] {replacement_0 = 0}
       : tensor<5xf32> into tensor<7xf32>
   return
 }
+
+// -----
+
+func.func @cast_like_insert_slice_dynamic(
+    %t: tensor<1x?x1xf32>, %f: f32, %pos: index) {
+  %c0 = arith.constant 0 : index
+  %0 = tensor.insert %f into %t[%c0, %pos, %c0] {replaced} : tensor<1x?x1xf32>
+
+  // Rank reduction
+  %c1 = arith.constant 1 : index
+  %dim1 = tensor.dim %t, %c1 : tensor<1x?x1xf32>
+  %1 = tensor.extract_slice %t[0, 0, 0][1, %dim1, 1][1, 1, 1]
+      : tensor<1x?x1xf32> to tensor<?xf32>
+  // expected-remark @below {{replacement found}}
+  %2 = tensor.insert %f into %1[%c0] : tensor<?xf32>
+  // 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
+}

utils/bazel/llvm-project-overlay/mlir/BUILD.bazel

cc_library(
     deps = [
         ":AffineDialect",
         ":IR",
         ":SCFDialect",
         ":TensorDialect",
         ":TensorTransformOpsIncGen",
         ":TensorTransforms",
         ":TransformDialect",
+        ":ValueBoundsOpInterface",
         "//llvm:Support",
     ],
 )
 
 cc_library(
     name = "Rewrite",
     srcs = glob([
         "lib/Rewrite/*.cpp",