[mlir][transform] Add TrackingListener

[springerm]

Summary

This change makes it possible to use a greedy pattern rewrite as part of a transform op, even if the transform op does not invalidate the target handle (in particular transform ops without FunctionalStyleTransformOpTrait) and the targeted op is not isolated from above.

The listener API allows us to track replacements of ops with values, but not ops with ops. Therefore, the TrackingListener is conservative: If an op is replaced with values that all have the same defining op and the defining op is of the same type as the original op, it is safe to assume that the op was replaced with an equivalent op. Otherwise, the op mapping is dropped. When this is not good enough, transforms can track values instead or provide a custom findReplacementOp function.

Depends On: D147038

[springerm]

Note: This listener can be used as a replacement for the TrackingListener in IREE. (It passes the CI.)

mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td

code commonExtraClassDeclaration = [{
     SmallVector<ReassociationIndices, 4> getReassociationIndices() {
       SmallVector<ReassociationIndices, 4> reassociationIndices;
       for (auto attr : getReassociation())
         reassociationIndices.push_back(llvm::to_vector<2>(
             llvm::map_range(attr.cast<ArrayAttr>(), [&](Attribute indexAttr) {
               return indexAttr.cast<IntegerAttr>().getInt();
             })));
       return reassociationIndices;
-    };
+    }
     RankedTensorType getSrcType() {
       return getSrc().getType().cast<RankedTensorType>();
     }
     RankedTensorType getResultType() {
       return getResult().getType().cast<RankedTensorType>();
     }
   }];
 

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

+//===- TensorTransformOps.h - Tensor transformation ops ---------*- C++ -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_DIALECT_TENSOR_TRANSFORMOPS_TENSORTRANSFORMOPS_H
+#define MLIR_DIALECT_TENSOR_TRANSFORMOPS_TENSORTRANSFORMOPS_H
+
+#include "mlir/Dialect/PDL/IR/PDLTypes.h"
+#include "mlir/Dialect/Transform/IR/TransformOps.h"
+#include "mlir/IR/OpImplementation.h"
+#include "mlir/IR/PatternMatch.h"
+
+namespace mlir {
+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:
+  explicit TrackingListener(transform::TransformState &state)
+      : transform::TrackingListener(state) {}
+
+protected:
+  Operation *findReplacementOp(Operation *op,
+                               ValueRange newValues) const override;
+};
+
+} // namespace tensor
+} // namespace mlir
+
+#endif // MLIR_DIALECT_TENSOR_TRANSFORMOPS_TENSORTRANSFORMOPS_H

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

 #define MLIR_DIALECT_TRANSFORM_IR_TRANSFORMOPS_H
 
 #include "mlir/Dialect/PDL/IR/PDLTypes.h"
 #include "mlir/Dialect/Transform/IR/TransformInterfaces.h"
 #include "mlir/Dialect/Transform/IR/TransformTypes.h"
 #include "mlir/IR/FunctionInterfaces.h"
 #include "mlir/IR/OpDefinition.h"
 #include "mlir/IR/OpImplementation.h"
+#include "mlir/IR/PatternMatch.h"
 #include "mlir/IR/SymbolTable.h"
 #include "mlir/Interfaces/CallInterfaces.h"
 #include "mlir/Interfaces/CastInterfaces.h"
 #include "mlir/Interfaces/ControlFlowInterfaces.h"
 
 namespace mlir {
 namespace transform {
 enum class FailurePropagationMode : uint32_t;
 class FailurePropagationModeAttr;
 
 /// A builder function that populates the body of a SequenceOp.
 using SequenceBodyBuilderFn = ::llvm::function_ref<void(
     ::mlir::OpBuilder &, ::mlir::Location, ::mlir::BlockArgument)>;
 using SequenceBodyBuilderArgsFn =
     ::llvm::function_ref<void(::mlir::OpBuilder &, ::mlir::Location,
                               ::mlir::BlockArgument, ::mlir::ValueRange)>;
+
+/// A listener that updates a TransformState based on IR modifications. This
+/// listener can be used during a greedy pattern rewrite to keep the transform
+/// state up-to-date.
+class TrackingListener : public RewriterBase::Listener,
+                         public TransformState::Extension {
+public:
+  explicit TrackingListener(TransformState &state)
+      : TransformState::Extension(state) {}
+
+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 newly-created op, which also has the same type as the given op,
+  /// that defining op is used as a replacement.
+  virtual Operation *findReplacementOp(Operation *op,
+                                       ValueRange newValues) const;
+
+  /// Return "true" if the given op is a new op.
+  bool isNewOp(Operation *op) const;
+
+  /// Return the single op that defines all given values (if any).
+  static Operation *getCommonDefiningOp(ValueRange values);
+
+private:
+  void notifyOperationInserted(Operation *op) override;
+
+  void notifyOperationRemoved(Operation *op) override;
+
+  void notifyOperationReplaced(Operation *op, ValueRange newValues) override;
+
+  /// Ops that were newly created during the transform.
+  DenseMap<OperationName, DenseSet<Operation *>> newOps;
+};
+
 } // namespace transform
 } // namespace mlir
 
 #define GET_OP_CLASSES
 #include "mlir/Dialect/Transform/IR/TransformOps.h.inc"
 
 #endif // MLIR_DIALECT_TRANSFORM_IR_TRANSFORMOPS_H

mlir/lib/Dialect/Tensor/CMakeLists.txt

 add_subdirectory(IR)
 add_subdirectory(Transforms)
+add_subdirectory(TransformOps)
 add_subdirectory(Utils)

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
+
+  LINK_LIBS PUBLIC
+  MLIRIR
+  MLIRTensorDialect
+  MLIRTransformDialect
+)

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/Tensor/IR/Tensor.h"
+#include "mlir/Dialect/Transform/IR/TransformDialect.h"
+#include "llvm/ADT/TypeSwitch.h"
+
+using namespace mlir;
+using namespace tensor;
+
+//===----------------------------------------------------------------------===//
+// TrackingListener
+//===----------------------------------------------------------------------===//
+
+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.
+    // 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()); })
+        .Default([](Operation *op) {});
+  } while (!values.empty());
+
+  return nullptr;
+}

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

if (succeeded(applicator.matchAndRewrite(op, rewriter)))
       results.push_back(op);
   });
 
   return success();
 }
 } // namespace
 
 //===----------------------------------------------------------------------===//
+// TrackingListener
+//===----------------------------------------------------------------------===//
+
+Operation *transform::TrackingListener::getCommonDefiningOp(ValueRange values) {
+  Operation *defOp = nullptr;
+  for (Value v : values) {
+    // Skip empty values.
+    if (!v)
+      continue;
+    if (!defOp) {
+      defOp = v.getDefiningOp();
+      continue;
+    }
+    if (defOp != v.getDefiningOp())
+      return nullptr;
+  }
+  return defOp;
+}
+
+Operation *
+transform::TrackingListener::findReplacementOp(Operation *op,
+                                               ValueRange newValues) const {
+  assert(op->getNumResults() == newValues.size() &&
+         "invalid number of replacement values");
+
+  // If the replacement values belong to different ops, drop the mapping.
+  Operation *defOp = getCommonDefiningOp(newValues);
+  if (!defOp)
+    return nullptr;
+
+  // If the replacement op has a different type, drop the mapping.
+  if (op->getName() != defOp->getName())
+    return nullptr;
+
+  // If the replacement op is not a new op, drop the mapping.
+  if (!isNewOp(defOp))
+    return nullptr;
+
+  return defOp;
+}
+
+bool transform::TrackingListener::isNewOp(Operation *op) const {
+  auto it = newOps.find(op->getName());
+  if (it == newOps.end())
+    return false;
+  return it->second.contains(op);
+}
+
+void transform::TrackingListener::notifyOperationInserted(Operation *op) {
+  newOps[op->getName()].insert(op);
+}
+
+void transform::TrackingListener::notifyOperationRemoved(Operation *op) {
+  // TODO: Walk can be removed when D144193 has landed.
+  op->walk([&](Operation *op) {
+    // Keep set of new ops up-to-date.
+    auto it = newOps.find(op->getName());
+    if (it != newOps.end())
+      it->second.erase(op);
+    // Remove mappings for result values.
+    for (OpResult value : op->getResults())
+      (void)replacePayloadValue(value, nullptr);
+    // Remove mapping for op.
+    (void)replacePayloadOp(op, nullptr);
+  });
+}
+
+void transform::TrackingListener::notifyOperationReplaced(
+    Operation *op, ValueRange newValues) {
+  assert(op->getNumResults() == newValues.size() &&
+         "invalid number of replacement values");
+  if (op->getNumResults() == 0)
+    return;
+
+  // Replace value handles.
+  for (auto [oldValue, newValue] : llvm::zip(op->getResults(), newValues))
+    (void)replacePayloadValue(oldValue, newValue);
+
+  // Replace op handle.
+  (void)replacePayloadOp(op, findReplacementOp(op, newValues));
+}
+
+//===----------------------------------------------------------------------===//
 // AlternativesOp
 //===----------------------------------------------------------------------===//
 
 OperandRange transform::AlternativesOp::getSuccessorEntryOperands(
     std::optional<unsigned> index) {
   if (index && getOperation()->getNumOperands() == 1)
     return getOperation()->getOperands();
   return OperandRange(getOperation()->operand_end(),

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

deps = [
         ":TilingInterface",
         ":Transforms",
         ":VectorDialect",
         "//llvm:Support",
     ],
 )
 
 cc_library(
+    name = "TensorTransformOps",
+    srcs = glob(["lib/Dialect/Tensor/TransformOps/*.cpp"]),
+    hdrs = glob(["include/mlir/Dialect/Tensor/TransformOps/*.h"]),
+    includes = ["include"],
+    deps = [
+        ":IR",
+        ":TensorDialect",
+        ":TransformDialect",
+        "//llvm:Support",
+    ],
+)
+
+cc_library(
     name = "Rewrite",
     srcs = glob([
         "lib/Rewrite/*.cpp",
         "lib/Rewrite/*.h",
     ]),
     hdrs = glob(["include/mlir/Rewrite/*.h"]),
     includes = ["include"],
     deps = [
deps = [
         ":ShapeTransforms",
         ":ShapeTransformsPassIncGen",
         ":SparseTensorDialect",
         ":SparseTensorPipelines",
         ":SparseTensorTransforms",
         ":TensorDialect",
         ":TensorInferTypeOpInterfaceImpl",
         ":TensorTilingInterfaceImpl",
+        ":TensorTransformOps",
         ":TensorTransforms",
         ":TosaDialect",
         ":TosaToLinalg",
         ":TransformDialect",
         ":TransformDialectTransforms",
         ":Transforms",
         ":TransformsPassIncGen",
         ":VectorDialect",