[mlir][transform] TrackingListener: Distinguish between failure and "should be dropped"

[springerm]

Summary

When looking for replacement ops (findReplacementOp) distinguish between "no replacement could be found" and "this op should be dropped from the mapping". The latter case will be utilized in a subsequent revision when a payload op is mapped to a consumed handle.

Depends On: D152374

[ftynse]

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

305–308

Nit: expand auto

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

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.
   ///
+  /// A "failure" return value indicates that no replacement operation could be
+  /// found. A "nullptr" return value indicates that no replacement op is needed
+  /// (e.g., handle is dead or was consumed) and that the payload op should
+  /// be dropped from the mapping.
+  ///
   /// 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
protected:
   /// 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,
+  virtual FailureOr<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,
                      function_ref<void(Diagnostic &)> reasonCallback) override;
 

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

if (!defOp) {
       continue;
     }
     if (defOp != v.getDefiningOp())
       return nullptr;
   }
   return defOp;
 }
 
-Operation *
+FailureOr<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(values);
     if (!defOp)
-      return nullptr;
+      return failure();
 
     // If the defining op has the same type, we take it as a replacement.
     if (op->getName() == defOp->getName())
       return defOp;
 
     // Replacing an op with a constant-like equivalent is a common
     // canonicalization.
     if (defOp->hasTrait<OpTrait::ConstantLike>())
do {
 
     // Skip through ops that implement CastOpInterface.
     if (isa<CastOpInterface>(defOp)) {
       values.assign(defOp->getOperands().begin(), defOp->getOperands().end());
       continue;
     }
   } while (!values.empty());
 
-  return nullptr;
+  return failure();
 }
 
 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";
auto hasAliveUser = [&]() {
     return false;
   };
   if (!hasAliveUser()) {
     // The op is tracked but the corresponding handles are dead.
     (void)replacePayloadOp(op, nullptr);
     return;
   }
 
-  Operation *replacement = findReplacementOp(op, newValues);
+  FailureOr<Operation *> replacement = findReplacementOp(op, newValues);
   // If the op is tracked but no replacement op was found, send a
   // notification.
-  if (!replacement)
+  if (failed(replacement)) {
     notifyPayloadReplacementNotFound(op, newValues);
-  (void)replacePayloadOp(op, replacement);
+    (void)replacePayloadOp(op, nullptr);
+    return;
+  }
+
+  (void)replacePayloadOp(op, *replacement);
 }
 
 transform::ErrorCheckingTrackingListener::~ErrorCheckingTrackingListener() {
   // The state of the ErrorCheckingTrackingListener must be checked and reset
   // if there was an error. This is to prevent errors from accidentally being
   // missed.
   assert(status.succeeded() && "listener state was not checked");
 }

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

 
 namespace {
 class DummyTrackingListener : public transform::TrackingListener {
 public:
   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);
+    FailureOr<Operation *> replacementOp = findReplacementOp(op, newValues);
+    if (failed(replacementOp))
+      return nullptr;
+    return *replacementOp;
   }
 };
 } // namespace
 
 static LogicalResult testTrackingListenerReplacements(Operation *rootOp) {
   // Find replaced op.
   Operation *replaced = nullptr;
   WalkResult status = rootOp->walk([&](Operation *op) {

mlir/test/lib/Dialect/Transform/TestTransformDialectExtension.cpp

 /// A TrackingListener for test cases. When the replacement op is
 /// "test.update_mapping", it is considered as a replacement op in the transform
 /// state mapping. Otherwise, it is not and the original op is simply removed
 /// from the mapping.
 class TestTrackingListener : public transform::TrackingListener {
   using transform::TrackingListener::TrackingListener;
 
 protected:
-  Operation *findReplacementOp(Operation *op,
-                               ValueRange newValues) const override {
+  FailureOr<Operation *>
+  findReplacementOp(Operation *op, ValueRange newValues) const override {
     if (newValues.size() != 1)
-      return nullptr;
+      return failure();
     Operation *replacement = newValues[0].getDefiningOp();
     if (!replacement)
-      return nullptr;
+      return failure();
     if (replacement->getName().getStringRef() != "test.update_mapping")
-      return nullptr;
+      return failure();
     return replacement;
   }
 };
 } // namespace
 
 DiagnosedSilenceableFailure
 mlir::test::TestTrackedRewriteOp::apply(transform::TransformResults &results,
                                         transform::TransformState &state) {