[mlir][transform] Utilize op interface instead of tensor::TrackingLister

[springerm]

Summary

Add a new interface FindPayloadReplacementOpInterface to specify ops that should be skipped when looking for payload replacement ops. Such ops are typically metadata-only ops.

With this change, we no longer need to maintain a custom TrackingListener in the tensor dialect.

Note: CastOpInterface by itself is not sufficient. Some metadata-only ops such as "tensor.reshape" are not casts, and it would be incorrect for them to implement the CastOpInterface.

Depends On: D146972

[nicolasvasilache]

Nice cleanup, thanks!

mlir/include/mlir/Dialect/Tensor/TransformOps/TensorTransformOps.h

 #include "mlir/Dialect/Transform/IR/TransformTypes.h"
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/IR/PatternMatch.h"
 
 namespace mlir {
 class DialectRegistry;
 
 namespace tensor {
-
-/// A specialized TrackingListener for transform ops that operate on tensor IR.
-/// This listener skips cast-like tensor ops when looking for payload op
-/// replacements.
-class TrackingListener : public transform::TrackingListener {
-public:
-  using transform::TrackingListener::TrackingListener;
-
-protected:
-  Operation *findReplacementOp(Operation *op,
-                               ValueRange newValues) const override;
-};
-
 void registerTransformDialectExtension(DialectRegistry &registry);
-
+void registerFindPayloadReplacementOpInterfaceExternalModels(
+    DialectRegistry &registry);
 } // namespace tensor
 } // namespace mlir
 
 #define GET_OP_CLASSES
 #include "mlir/Dialect/Tensor/TransformOps/TensorTransformOps.h.inc"
 
 #endif // MLIR_DIALECT_TENSOR_TRANSFORMOPS_TENSORTRANSFORMOPS_H

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

 def NavigationTransformOpTrait : NativeOpTrait<"NavigationTransformOpTrait"> {
   let cppNamespace = "::mlir::transform";
 }
 
 def ParamProducerTransformOpTrait : NativeOpTrait<"ParamProducerTransformOpTrait"> {
   let cppNamespace = "::mlir::transform";
 }
 
+def FindPayloadReplacementOpInterface
+    : OpInterface<"FindPayloadReplacementOpInterface"> {
+  let description = [{
+    This interface is queried by the `TrackingListener` and can be implemented
+    by payload ops to indicate that the lookup should be continue with its
+    operands when looking for payload op replacements.
+
+    Example: Consider the case where a tracked "test.foo" payload op is replaced
+    with a new "test.foo" op, but wrapped in a "tensor.reshape" op. In that
+    case, the mapping of the original "test.foo" op should be updated with the
+    new "test.foo" op. A "tensor.reshape" is a metadata-only op that should be
+    skipped when inspecting the replacement values of the original "test.foo"
+    op. More details can be found at `TrackingListener` documentation.
+
+    Note: Ops that implement `CastOpInterface` do not need to implement this
+    interface. Such ops are skipped by default. This interface should be
+    implemented by cast-like/metadata-only ops that cannot implement
+    `CastOpInterface`.
+  }];
+
+  let cppNamespace = "::mlir::transform";
+
+  let methods = [
+    InterfaceMethod<
+      /*desc=*/[{
+        Return the operands at which the lookup for replacement payload ops
+        should continue.
+      }],
+      /*returnType=*/"::llvm::SmallVector<::mlir::Value>",
+      /*name=*/"getNextOperands",
+      /*arguments=*/(ins)
+    >,
+  ];
+}
+
 #endif // MLIR_DIALECT_TRANSFORM_IR_TRANSFORM_INTERFACES_TD

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

public:
   explicit TrackingListener(TransformState &state, TransformOpInterface op)
       : TransformState::Extension(state), transformOp(op) {}
 
 protected:
   /// Return a replacement payload op for the given op, which is going to be
   /// replaced with the given values. By default, if all values are defined by
   /// the same op, which also has the same type as the given op, that defining
   /// op is used as a replacement.
+  ///
+  /// Example: A tracked "linalg.generic" with two results is replaced with two
+  /// values defined by (another) "linalg.generic". It is reasonable to assume
+  /// that the replacement "linalg.generic" represents the same "computation".
+  /// Therefore, the payload op mapping is updated to the defining op of the
+  /// replacement values.
+  ///
+  /// Counter Example: A "linalg.generic" is replaced with values defined by an
+  /// "scf.for". Without further investigation, the relationship between the
+  /// "linalg.generic" and the "scf.for" is unclear. They may not represent the
+  /// same computation; e.g., there may be tiled "linalg.generic" inside the
+  /// loop body that represents the original computation. Therefore, the
+  /// TrackingListener is conservative by default: it drops the mapping and
+  /// triggers the "payload replacement not found" notification.
+  ///
+  /// If no replacement op could be found according to the rules mentioned
+  /// above, this function tries to skip over cast-like ops that implement
+  /// `CastOpInterface`.
+  ///
+  /// Example: A tracked "linalg.generic" is replaced with "linalg.generic",
+  /// wrapped in a "tensor.cast". A cast is a metadata-only operation and it is
+  /// reasonable to assume that the wrapped "linalg.generic" represents the same
+  /// computation as the original "linalg.generic". The mapping is updated
+  /// accordingly.
+  ///
+  /// Certain ops (typically also metadata-only ops) are not considered casts,
+  /// but should be skipped nonetheless. Such ops should implement
+  /// `FindPayloadReplacementOpInterface` to specify with which operands the
+  /// lookup should continue.
+  ///
+  /// Example: A tracked "linalg.generic" is replaced with "linalg.generic",
+  /// wrapped in a "tensor.reshape". A reshape is a metadata-only operation but
+  /// not cast. (Implementing `CastOpInterface` would be incorrect and cause
+  /// invalid foldings.) However, due to its `FindPayloadReplacementOpInterface`
+  /// implementation, the replacement op lookup continues with the wrapped
+  /// "linalg.generic" and the mapping is updated accordingly.
+  ///
+  /// Derived classes may override `findReplacementOp` to specify custom
+  /// replacement rules.
   virtual Operation *findReplacementOp(Operation *op,
                                        ValueRange newValues) const;
 
   /// Notify the listener that the pattern failed to match the given operation,
   /// and provide a callback to populate a diagnostic with the reason why the
   /// failure occurred.
   LogicalResult
   notifyMatchFailure(Location loc,

mlir/include/mlir/InitAllDialects.h

inline void registerAllDialects(DialectRegistry &registry) {
   memref::registerRuntimeVerifiableOpInterfaceExternalModels(registry);
   memref::registerValueBoundsOpInterfaceExternalModels(registry);
   memref::registerMemorySlotExternalModels(registry);
   scf::registerBufferizableOpInterfaceExternalModels(registry);
   scf::registerValueBoundsOpInterfaceExternalModels(registry);
   shape::registerBufferizableOpInterfaceExternalModels(registry);
   sparse_tensor::registerBufferizableOpInterfaceExternalModels(registry);
   tensor::registerBufferizableOpInterfaceExternalModels(registry);
+  tensor::registerFindPayloadReplacementOpInterfaceExternalModels(registry);
   tensor::registerInferTypeOpInterfaceExternalModels(registry);
   tensor::registerTilingInterfaceExternalModels(registry);
   tensor::registerValueBoundsOpInterfaceExternalModels(registry);
   vector::registerBufferizableOpInterfaceExternalModels(registry);
 }
 
 /// Append all the MLIR dialects to the registry contained in the given context.
 inline void registerAllDialects(MLIRContext &context) {
   DialectRegistry registry;
   registerAllDialects(registry);
   context.appendDialectRegistry(registry);
 }
 
 } // namespace mlir
 
 #endif // MLIR_INITALLDIALECTS_H_

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

 
 #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/Tensor/Utils/Utils.h"
 #include "mlir/Dialect/Transform/IR/TransformDialect.h"
 #include "mlir/Dialect/Transform/IR/TransformInterfaces.h"
-#include "llvm/ADT/TypeSwitch.h"
 
 using namespace mlir;
 using namespace tensor;
 
 //===----------------------------------------------------------------------===//
-// TrackingListener
+// FindPayloadReplacementOpInterface implementations
 //===----------------------------------------------------------------------===//
 
-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;
-
-    Operation *defOp = getCommonDefiningOp(values);
-    if (!defOp)
-      return nullptr;
-
-    // Skip cast-like operations.
-    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());
+namespace {
+struct ExtractSliceOpReplacementInterface
+    : public transform::FindPayloadReplacementOpInterface::ExternalModel<
+          ExtractSliceOpReplacementInterface, tensor::ExtractSliceOp> {
+  SmallVector<Value> getNextOperands(Operation *op) const {
+    auto extractSliceOp = cast<tensor::ExtractSliceOp>(op);
+    if (!isCastLikeExtractSliceOp(extractSliceOp))
+      return {};
+    return {extractSliceOp.getSource()};
+  }
+};
+
+struct InsertSliceOpReplacementInterface
+    : public transform::FindPayloadReplacementOpInterface::ExternalModel<
+          InsertSliceOpReplacementInterface, tensor::InsertSliceOp> {
+  SmallVector<Value> getNextOperands(Operation *op) const {
+    auto insertSliceOp = cast<tensor::InsertSliceOp>(op);
+    if (!isCastLikeInsertSliceOp(insertSliceOp))
+      return {};
+    return {insertSliceOp.getSource()};
+  }
+};
+
+struct ReshapeOpReplacementInterface
+    : public transform::FindPayloadReplacementOpInterface::ExternalModel<
+          ReshapeOpReplacementInterface, tensor::ReshapeOp> {
+  SmallVector<Value> getNextOperands(Operation *op) const {
+    auto reshapeOp = cast<tensor::ReshapeOp>(op);
+    return {reshapeOp.getSource()};
+  }
+};
 
-  return nullptr;
+template <typename ConcreteOp>
+struct ReassociativeReshapeOpReplacementInterface
+    : public transform::FindPayloadReplacementOpInterface::ExternalModel<
+          ReassociativeReshapeOpReplacementInterface<ConcreteOp>, ConcreteOp> {
+  SmallVector<Value> getNextOperands(Operation *op) const {
+    auto reshapeOp = cast<ConcreteOp>(op);
+    return {reshapeOp.getSrc()};
+  }
+};
+} // namespace
+
+void tensor::registerFindPayloadReplacementOpInterfaceExternalModels(
+    DialectRegistry &registry) {
+  registry.addExtension(+[](MLIRContext *ctx, tensor::TensorDialect *dialect) {
+    CollapseShapeOp::attachInterface<
+        ReassociativeReshapeOpReplacementInterface<CollapseShapeOp>>(*ctx);
+    ExpandShapeOp::attachInterface<
+        ReassociativeReshapeOpReplacementInterface<ExpandShapeOp>>(*ctx);
+    ExtractSliceOp::attachInterface<ExtractSliceOpReplacementInterface>(*ctx);
+    InsertSliceOp::attachInterface<InsertSliceOpReplacementInterface>(*ctx);
+    ReshapeOp::attachInterface<ReshapeOpReplacementInterface>(*ctx);
+  });
 }
 
 //===----------------------------------------------------------------------===//
 // MakeLoopIndependentOp
 //===----------------------------------------------------------------------===//
 
 DiagnosedSilenceableFailure transform::MakeLoopIndependentOp::applyToOne(
     Operation *target, transform::ApplyToEachResultList &results,

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

Operation *transform::TrackingListener::getCommonDefiningOp(ValueRange values) {
   return defOp;
 }
 
 Operation *
 transform::TrackingListener::findReplacementOp(Operation *op,
                                                ValueRange newValues) const {
   assert(op->getNumResults() == newValues.size() &&
          "invalid number of replacement values");
+  SmallVector<Value> values(newValues.begin(), newValues.end());
 
+  do {
     // If the replacement values belong to different ops, drop the mapping.
-  Operation *defOp = getCommonDefiningOp(newValues);
+    Operation *defOp = getCommonDefiningOp(values);
     if (!defOp)
       return nullptr;
 
-  // If the replacement op has a different type, drop the mapping.
-  if (op->getName() != defOp->getName())
-    return nullptr;
-
+    // If the defining op has the same type, we take it as a replacement.
+    if (op->getName() == defOp->getName())
       return defOp;
+
+    values.clear();
+
+    // Skip through ops that implement FindPayloadReplacementOpInterface.
+    if (auto findReplacementOpInterface =
+            dyn_cast<FindPayloadReplacementOpInterface>(defOp)) {
+      values.assign(findReplacementOpInterface.getNextOperands());
+      continue;
+    }
+
+    // Skip through ops that implement CastOpInterface.
+    if (isa<CastOpInterface>(defOp)) {
+      values.assign(defOp->getOperands().begin(), defOp->getOperands().end());
+      continue;
+    }
+  } while (!values.empty());
+
+  return nullptr;
 }
 
 LogicalResult transform::TrackingListener::notifyMatchFailure(
     Location loc, function_ref<void(Diagnostic &)> reasonCallback) {
   LLVM_DEBUG({
     Diagnostic diag(loc, DiagnosticSeverity::Remark);
     reasonCallback(diag);
     DBGS() << "Match Failure : " << diag.str() << "\n";

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/Dialect/Transform/IR/TransformOps.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 
 using namespace mlir;
 
 namespace {
 struct TestTensorTransforms
     : public PassWrapper<TestTensorTransforms, OperationPass<>> {
patterns.add<RewriteExtractSliceFromCollapseShapeUsingScfForeach>(
         rootOp->getContext());
   else
     patterns.add<RewriteExtractSliceFromCollapseShapeUsingScfFor>(
         rootOp->getContext());
   return applyPatternsAndFoldGreedily(rootOp, std::move(patterns));
 }
 
 namespace {
-class DummyTrackingListener : public tensor::TrackingListener {
+class DummyTrackingListener : public transform::TrackingListener {
 public:
-  using tensor::TrackingListener::TrackingListener;
+  using transform::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
 

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",