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

[springerm]

Summary

Certain InsertSliceOps, that do not use elements from the destination, are treated like casts when looking for replacement ops. Such InsertSliceOps are typically rank expansions.

Tensors with dynamic shape are not supported at the moment.

Also adds test cases for the TrackingListener.

mlir/include/mlir/Dialect/Transform/IR/TransformInterfaces.h

 void prepareValueMappings(
     SmallVectorImpl<SmallVector<transform::MappedValue>> &mappings,
     ValueRange values, const transform::TransformState &state);
 
 /// Populates `results` with payload associations that match exactly those of
 /// the operands to `block`'s terminator.
 void forwardTerminatorOperands(Block *block, transform::TransformState &state,
                                transform::TransformResults &results);
+
+/// Make a dummy transform state for testing purposes. This MUST NOT be used
+/// outside of test cases.
+TransformState makeTransformStateForTesting(Region *region,
+                                            Operation *payloadRoot);
 } // namespace detail
 
 /// Options controlling the application of transform operations by the
 /// TransformState.
 class TransformOptions {
 public:
   TransformOptions() = default;
   TransformOptions(const TransformOptions &) = default;
struct Mappings {
     ValueMapping values;
     ValueMapping reverseValues;
   };
 
   friend LogicalResult applyTransforms(Operation *, TransformOpInterface,
                                        const RaggedArray<MappedValue> &,
                                        const TransformOptions &);
 
+  friend TransformState
+  detail::makeTransformStateForTesting(Region *region, Operation *payloadRoot);
+
 public:
   /// Returns the op at which the transformation state is rooted. This is
   /// typically helpful for transformations that apply globally.
   Operation *getTopLevel() const;
 
   /// Returns the number of extra mappings for the top-level operation.
   size_t getNumTopLevelMappings() const { return topLevelMappedValues.size(); }
 

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

 
 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.
+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())
+    return false;
+  return op.getSourceType().getNumElements() ==
+         op.getDestType().getNumElements();
+}
+
 Operation *
 tensor::TrackingListener::findReplacementOp(Operation *op,
                                             ValueRange newValues) const {
   SmallVector<Value> values(newValues.begin(), newValues.end());
   do {
     if (Operation *replacement =
             transform::TrackingListener::findReplacementOp(op, values))
       return replacement;
do {
     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());
+        })
         .Default([](Operation *op) {});
   } while (!values.empty());
 
   return nullptr;
 }
 
 //===----------------------------------------------------------------------===//
 // MakeLoopIndependentOp

mlir/lib/Dialect/Transform/IR/TransformInterfaces.cpp

if (llvm::isa<transform::TransformHandleTypeInterface>(result.getType())) {
       assert(
           llvm::isa<transform::TransformParamTypeInterface>(result.getType()) &&
           "unhandled transform type interface");
       results.setParams(result, state.getParams(terminatorOperand));
     }
   }
 }
 
+transform::TransformState
+transform::detail::makeTransformStateForTesting(Region *region,
+                                                Operation *payloadRoot) {
+  return TransformState(region, payloadRoot);
+}
+
 //===----------------------------------------------------------------------===//
 // Utilities for PossibleTopLevelTransformOpTrait.
 //===----------------------------------------------------------------------===//
 
 /// Appends to `effects` the memory effect instances on `target` with the same
 /// resource and effect as the ones the operation `iface` having on `source`.
 static void
 remapEffects(MemoryEffectOpInterface iface, BlockArgument source, Value target,

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

+// RUN: mlir-opt -test-tensor-transform-patterns=test-tracking-listener \
+// RUN:     -split-input-file -verify-diagnostics %s
+
+func.func @replace_op_with_op_of_same_type() {
+  %0 = "test.foo"() {replaced} : () -> (tensor<5xf32>)
+  // expected-remark @below {{replacement found}}
+  %1 = "test.foo"() {replacement_0 = 0} : () -> (tensor<5xf32>)
+  return
+}
+
+// -----
+
+func.func @replace_op_with_op_of_different_type() {
+  // expected-error @below {{listener could not find replacement op}}
+  %0 = tensor.empty() {replaced} : tensor<5xf32>
+  %1 = "test.foo"() {replacement_0 = 0} : () -> (tensor<5xf32>)
+  return
+}
+
+// -----
+
+func.func @multi_result_replacement() {
+  %0:2 = "test.foo"() {replaced} : () -> (tensor<5xf32>, tensor<6xf32>)
+  // expected-remark @below {{replacement found}}
+  %1:2 = "test.foo"() {replacement_0 = 0, replacement_1 = 1}
+      : () -> (tensor<5xf32>, tensor<6xf32>)
+  return
+}
+
+// -----
+
+func.func @multi_result_replacement_with_multiple_ops() {
+  // expected-error @below {{listener could not find replacement op}}
+  %0:2 = "test.foo"() {replaced} : () -> (tensor<5xf32>, tensor<6xf32>)
+  %1:2 = "test.foo"() {replacement_0 = 0} : () -> (tensor<5xf32>, tensor<6xf32>)
+  %2:2 = "test.foo"() {replacement_1 = 1} : () -> (tensor<5xf32>, tensor<6xf32>)
+  return
+}
+
+// -----
+
+func.func @replacement_wrapped_in_cast() {
+  %0 = "test.foo"() {replaced} : () -> (tensor<5xf32>)
+  // expected-remark @below {{replacement found}}
+  %1 = "test.foo"() : () -> (tensor<?xf32>)
+  %2 = tensor.cast %1 {replacement_0 = 0} : tensor<?xf32> to tensor<5xf32>
+  return
+}
+
+// -----
+
+func.func @replacement_wrapped_in_chain_of_casts() {
+  %0 = "test.foo"() {replaced} : () -> (tensor<5xf32>)
+  // expected-remark @below {{replacement found}}
+  %1 = "test.foo"() : () -> (tensor<?xf32>)
+  %2 = tensor.cast %1 : tensor<?xf32> to tensor<5xf32>
+  %3 = tensor.cast %2 : tensor<5xf32> to tensor<?xf32>
+  %4 = tensor.cast %3 {replacement_0 = 0} : tensor<?xf32> to tensor<5xf32>
+  return
+}
+
+// -----
+
+func.func @cast_like_insert_slice(%t: tensor<1x5xf32>) {
+  %0 = "test.foo"() {replaced} : () -> (tensor<5xf32>)
+  // expected-remark @below {{replacement found}}
+  %1 = "test.foo"() : () -> (tensor<5xf32>)
+  %2 = tensor.insert_slice %1 into %t[0, 0][1, 5][1, 1] {replacement_0 = 0}
+      : tensor<5xf32> into tensor<1x5xf32>
+  return
+}
+
+// -----
+
+func.func @non_cast_like_insert_slice(%t: tensor<7xf32>) {
+  // expected-error @below {{listener could not find replacement op}}
+  %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
+}

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

 # Exclude tests from libMLIR.so
 add_mlir_library(MLIRTensorTestPasses
   TestTensorTransforms.cpp
 
   EXCLUDE_FROM_LIBMLIR
 
   LINK_LIBS PUBLIC
   MLIRArithDialect
   MLIRLinalgDialect
   MLIRPass
   MLIRSCFDialect
   MLIRTensorDialect
+  MLIRTensorTransformOps
   MLIRTensorTransforms
+  MLIRTransformDialect
   MLIRTransforms
   )

mlir/test/lib/Dialect/Tensor/TestTensorTransforms.cpp

 //===- TestTensorTransforms.cpp - Test Tensor transformation patterns -----===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements logic for testing Tensor transformations.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/Dialect/Tensor/TransformOps/TensorTransformOps.h"
 #include "mlir/Dialect/Tensor/Transforms/TransformUtils.h"
 #include "mlir/Dialect/Tensor/Transforms/Transforms.h"
+#include "mlir/Dialect/Transform/IR/TransformInterfaces.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 
 using namespace mlir;
 
 namespace {
 struct TestTensorTransforms
     : public PassWrapper<TestTensorTransforms, OperationPass<>> {
Option<bool> useForeach{
           "Use the scf.forall operation when generating loop nests for "
           "the extract_slice of collapse_shape pattern"),
       llvm::cl::init(false)};
 
   Option<bool> testSimplifyPackPatterns{
       *this, "test-simplify-pack-patterns",
       llvm::cl::desc("Test patterns to simplify tensor.pack"),
       llvm::cl::init(false)};
+
+  Option<bool> testTrackingListener{
+      *this, "test-tracking-listener",
+      llvm::cl::desc("Test tensor TrackingListener for the transform dialect"),
+      llvm::cl::init(false)};
 };
 } // namespace
 
 static void applyReassociativeReshapeFoldingPatterns(Operation *rootOp) {
   RewritePatternSet patterns(rootOp->getContext());
   tensor::populateReassociativeReshapeFoldingPatterns(patterns);
   (void)applyPatternsAndFoldGreedily(rootOp, std::move(patterns));
 }
if (useForeach)
     patterns.add<RewriteExtractSliceFromCollapseShapeUsingScfForeach>(
         rootOp->getContext());
   else
     patterns.add<RewriteExtractSliceFromCollapseShapeUsingScfFor>(
         rootOp->getContext());
   return applyPatternsAndFoldGreedily(rootOp, std::move(patterns));
 }
 
+namespace {
+class DummyTrackingListener : public tensor::TrackingListener {
+public:
+  using tensor::TrackingListener::TrackingListener;
+
+  // Expose `findReplacementOp` as a public function, so that it can be tested.
+  Operation *getReplacementOp(Operation *op, ValueRange newValues) const {
+    return findReplacementOp(op, newValues);
+  }
+};
+} // namespace
+
+static LogicalResult testTrackingListenerReplacements(Operation *rootOp) {
+  // Find replaced op.
+  Operation *replaced = nullptr;
+  WalkResult status = rootOp->walk([&](Operation *op) {
+    if (op->hasAttr("replaced")) {
+      if (replaced) {
+        op->emitError("only one 'replaced' op is allowed per test case");
+        replaced->emitRemark("other 'replaced' op");
+        return WalkResult::interrupt();
+      }
+      replaced = op;
+    }
+    return WalkResult::advance();
+  });
+  if (status.wasInterrupted())
+    return failure();
+  if (!replaced) {
+    replaced->emitError("could not find 'replaced' op");
+    return failure();
+  }
+
+  // Find replacements.
+  SmallVector<Value> replacements(replaced->getNumResults(), Value());
+  status = rootOp->walk([&](Operation *op) {
+    for (int64_t i = 0; i < replaced->getNumResults(); ++i) {
+      if (auto attr = op->getAttrOfType<IntegerAttr>("replacement_" +
+                                                     std::to_string(i))) {
+        if (replacements[i]) {
+          op->emitError("only one 'replacement_" + std::to_string(i) +
+                        "' is allowed per test case");
+          replacements[i].getDefiningOp()->emitRemark("other 'replacement_" +
+                                                      std::to_string(i) + "'");
+          return WalkResult::interrupt();
+        }
+        replacements[i] = op->getResult(attr.getInt());
+      }
+    }
+    return WalkResult::advance();
+  });
+  if (status.wasInterrupted())
+    return failure();
+
+  if (!llvm::all_of(replacements,
+                    [](Value v) { return static_cast<bool>(v); })) {
+    replaced->emitError("insufficient replacement values");
+    return failure();
+  }
+
+  // Find the replacement op (if any) and emit a remark/error.
+  transform::TransformState transformState =
+      transform::detail::makeTransformStateForTesting(/*region=*/nullptr,
+                                                      /*payloadRoot=*/nullptr);
+  DummyTrackingListener listener(transformState,
+                                 transform::TransformOpInterface());
+  Operation *replacement = listener.getReplacementOp(replaced, replacements);
+  if (!replacement) {
+    replaced->emitError("listener could not find replacement op");
+    return failure();
+  }
+
+  replacement->emitRemark("replacement found");
+  return success();
+}
+
 void TestTensorTransforms::runOnOperation() {
   Operation *rootOp = getOperation();
   if (testSimplifyPackPatterns)
     applySimplifyPackPatterns(rootOp);
   if (testFoldConstantExtractSlice)
     applyFoldConstantExtractSlicePatterns(rootOp);
   if (testFoldConsecutiveInsertExtractSlice)
     applyFoldConsecutiveInsertExtractSlicePatterns(rootOp);
   if (testReassociativeReshapeFolding)
     applyReassociativeReshapeFoldingPatterns(rootOp);
   if (testEmptyOpFolding)
     applyEmptyOpFoldingPatterns(rootOp);
   if (testFoldIntoPackAndUnpack)
     applyFoldIntoPackAndUnpackPatterns(rootOp);
   if (testRewriteExtractSliceWithTiledCollapseShape) {
     if (failed(
             applyRewriteExtractFromCollapseShapePatterns(rootOp, useForeach)))
       return signalPassFailure();
   }
+  if (testTrackingListener)
+    if (failed(testTrackingListenerReplacements(rootOp)))
+      return signalPassFailure();
 }
 
 namespace mlir {
 namespace test {
 void registerTestTensorTransforms() {
   PassRegistration<TestTensorTransforms>();
 }
 } // namespace test
 } // namespace mlir

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

cc_library(
     defines = ["MLIR_CUDA_CONVERSIONS_ENABLED"],
     includes = ["lib/Dialect/Test"],
     deps = [
         "//mlir:ArithDialect",
         "//mlir:LinalgDialect",
         "//mlir:Pass",
         "//mlir:SCFDialect",
         "//mlir:TensorDialect",
+        "//mlir:TensorTransformOps",
         "//mlir:TensorTransforms",
+        "//mlir:TransformDialect",
         "//mlir:Transforms",
     ],
 )
 
 cc_library(
     name = "TestVector",
     srcs = glob(["lib/Dialect/Vector/*.cpp"]),
     defines = ["MLIR_CUDA_CONVERSIONS_ENABLED"],