[mlir][transform] Remove redundant handle check in `replacePayload...`

[springerm]

Summary

Depends On: D153765

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

#endif // NDEBUG
   /// as well as any association between value handles and the results of said
   /// payload IR op.
   void forgetMapping(Value opHandle, ValueRange origOpFlatResults);
 
   void forgetValueMapping(Value valueHandle,
                           ArrayRef<Operation *> payloadOperations);
 
   /// Replaces the given payload op with another op. If the replacement op is
-  /// null, removes the association of the payload op with its handle.
+  /// null, removes the association of the payload op with its handle. Returns
+  /// failure if the op is not associated with any handle.
   ///
   /// Note: This function does not update value handles. None of the original
   /// op's results are allowed to be mapped to any value handle.
   LogicalResult replacePayloadOp(Operation *op, Operation *replacement);
 
   /// Replaces the given payload value with another value. If the replacement
   /// value is null, removes the association of the payload value with its
-  /// handle.
+  /// handle. Returns failure if the value is not associated with any handle.
   LogicalResult replacePayloadValue(Value value, Value replacement);
 
   /// Records handle invalidation reporters into `newlyInvalidated`.
   /// Specifically,
   ///  - `handle` is the op operand that consumes the handle,
   ///  - `potentialAncestors` is a list of ancestors of the payload operation
   ///     that the consumed handle is associated with, including itself,
   ///  - `throughValue` is the payload value the handle to which is consumed,

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

for (Value opHandle : opHandles) {
       dropMappingEntry(localMappings.reverse, payloadOp, opHandle);
     }
   }
 }
 
 LogicalResult
 transform::TransformState::replacePayloadOp(Operation *op,
                                             Operation *replacement) {
-  // Drop the mapping between the op and all handles that point to it. Don't
-  // care if there are on such handles.
-  SmallVector<Value> opHandles;
-  (void)getHandlesForPayloadOp(op, opHandles);
-  for (Value handle : opHandles) {
-    Mappings &mappings = getMapping(handle);
-    dropMappingEntry(mappings.reverse, op, handle);
-  }
-
 #ifndef NDEBUG
   for (Value opResult : op->getResults()) {
     SmallVector<Value> valueHandles;
     (void)getHandlesForPayloadValue(opResult, valueHandles);
     assert(valueHandles.empty() && "expected no mapping to old results");
   }
 #endif // NDEBUG
 
 #if LLVM_ENABLE_ABI_BREAKING_CHECKS
   if (options.getExpensiveChecksEnabled()) {
     auto it = cachedNames.find(op);
     assert(it != cachedNames.end() && "entry not found");
     assert(it->second == op->getName() && "operation name mismatch");
     cachedNames.erase(it);
   }
 #endif // LLVM_ENABLE_ABI_BREAKING_CHECKS
 
+  // Drop the mapping between the op and all handles that point to it. Fail if
+  // there are no handles.
+  SmallVector<Value> opHandles;
+  if (failed(getHandlesForPayloadOp(op, opHandles)))
+    return failure();
+  for (Value handle : opHandles) {
+    Mappings &mappings = getMapping(handle);
+    dropMappingEntry(mappings.reverse, op, handle);
+  }
+
   // TODO: consider invalidating the handles to nested objects here.
 
 #if LLVM_ENABLE_ABI_BREAKING_CHECKS
   if (replacement && options.getExpensiveChecksEnabled()) {
     auto insertion = cachedNames.insert({replacement, replacement->getName()});
     if (!insertion.second) {
       assert(insertion.first->second == replacement->getName() &&
              "operation is already cached with a different name");
#endif // LLVM_ENABLE_ABI_BREAKING_CHECKS
   }
 
   return success();
 }
 
 LogicalResult
 transform::TransformState::replacePayloadValue(Value value, Value replacement) {
   SmallVector<Value> valueHandles;
-  (void)getHandlesForPayloadValue(value, valueHandles);
+  if (failed(getHandlesForPayloadValue(value, valueHandles)))
+    return failure();
 
   for (Value handle : valueHandles) {
     Mappings &mappings = getMapping(handle);
     dropMappingEntry(mappings.reverseValues, value, handle);
 
     // If replacing with null, that is erasing the mapping, drop the mapping
     // between the handles and the IR objects
     if (!replacement) {

 // TransformState::Extension
 //===----------------------------------------------------------------------===//
 
 transform::TransformState::Extension::~Extension() = default;
 
 LogicalResult
 transform::TransformState::Extension::replacePayloadOp(Operation *op,
                                                        Operation *replacement) {
-  SmallVector<Value> handles;
-  if (failed(state.getHandlesForPayloadOp(op, handles)))
-    return failure();
-
   // TODO: we may need to invalidate handles to operations and values nested in
   // the operation being replaced.
   return state.replacePayloadOp(op, replacement);
 }
 
 LogicalResult
 transform::TransformState::Extension::replacePayloadValue(Value value,
                                                           Value replacement) {
-  SmallVector<Value> handles;
-  if (failed(state.getHandlesForPayloadValue(value, handles)))
-    return failure();
-
   return state.replacePayloadValue(value, replacement);
 }
 
 //===----------------------------------------------------------------------===//
 // TransformState::RegionScope
 //===----------------------------------------------------------------------===//
 
 transform::TransformState::RegionScope::~RegionScope() {