[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);
-  }
+  // TODO: consider invalidating the handles to nested objects here.
 
 #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
 
+  // 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);
+  }
+
 #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
-
-  // 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 (replacement) {
+      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
 
   // Replace the pointed-to object of all handles with the replacement object.
   // In case a payload op was erased (replacement object is nullptr), a nullptr
   // is stored in the mapping. These nullptrs are removed after each transform.
   // Furthermore, nullptrs are not enumerated by payload op iterators. The
   // relative order of ops is preserved.
   //
#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) {
for (Operation *ancestor : potentialAncestors) {
     if (!ancestor->isAncestor(definingOp))
       continue;
 
     Operation *owner = opHandle.getOwner();
     unsigned operandNo = opHandle.getOperandNumber();
     Location ancestorLoc = ancestor->getLoc();
     Location opLoc = definingOp->getLoc();
     Location valueLoc = payloadValue.getLoc();
-    newlyInvalidated[valueHandle] =
-        [valueHandle, owner, operandNo, resultNo, argumentNo, blockNo, regionNo,
-         ancestorLoc, opLoc, valueLoc](Location currentLoc) {
-          InFlightDiagnostic diag = emitError(currentLoc)
-                                    << "op uses a handle invalidated by a "
-                                       "previously executed transform op";
-          diag.attachNote(valueHandle.getLoc()) << "invalidated handle";
-          diag.attachNote(owner->getLoc())
-              << "invalidated by this transform op that consumes its operand #"
-              << operandNo
-              << " and invalidates all handles to payload IR entities "
-                 "associated with this operand and entities nested in them";
-          diag.attachNote(ancestorLoc)
-              << "ancestor op associated with the consumed handle";
-          if (resultNo) {
-            diag.attachNote(opLoc)
-                << "op defining the value as result #" << *resultNo;
-          } else {
-            diag.attachNote(opLoc)
-                << "op defining the value as block argument #" << argumentNo
-                << " of block #" << blockNo << " in region #" << regionNo;
-          }
-          diag.attachNote(valueLoc) << "payload value";
-        };
+    newlyInvalidated[valueHandle] = [valueHandle, owner, operandNo, resultNo,
+                                     argumentNo, blockNo, regionNo, ancestorLoc,
+                                     opLoc, valueLoc](Location currentLoc) {
   }
 }
 
 void transform::TransformState::recordOpHandleInvalidation(
     OpOperand &handle, ArrayRef<Operation *> potentialAncestors,
     Value throughValue,
     transform::TransformState::InvalidatedHandleMap &newlyInvalidated) const {
 

 // 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() {