[mlir] improve the error message in expensive transform checks

[ftynse]

Summary

Include the transform op being applied when reporting it using an invalidated
handle. This was missing previously and made it harder for the user to
understand where the handle is being used, especially if the transform script
included some sort of iteration.

[guraypp]

looks good to me. It is a useful feedback for the user to understand what is going on in the conversion dialect.

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

private:
   /// The top-level operation that contains all payload IR, typically a module.
   Operation *topLevel;
 
   /// Additional options controlling the transformation state behavior.
   TransformOptions options;
 
   /// The mapping from invalidated handles to the error-reporting functions that
   /// describe when the handles were invalidated. Calling such a function emits
-  /// a user-visible diagnostic.
-  DenseMap<Value, std::function<void()>> invalidatedHandles;
+  /// a user-visible diagnostic with an additional note pointing to the given
+  /// location.
+  DenseMap<Value, std::function<void(Location)>> invalidatedHandles;
 
 #if LLVM_ENABLE_ABI_BREAKING_CHECKS
   /// A stack of nested regions that are being processed in the transform IR.
   /// Each region must be an ancestor of the following regions in this list.
   /// These are also the keys for "mappings".
   SmallVector<Region *> regionStack;
 #endif // LLVM_ENABLE_ABI_BREAKING_CHECKS
 };

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

for (const auto &kvp : mapping.reverse) {
         // by-reference because it will go out of scope. Additionally, extract
         // location from Payload IR ops because the ops themselves may be
         // deleted before the lambda gets called.
         Location ancestorLoc = ancestor->getLoc();
         Location opLoc = op->getLoc();
         Operation *owner = handle.getOwner();
         unsigned operandNo = handle.getOperandNumber();
         invalidatedHandles[otherHandle] = [ancestorLoc, opLoc, owner, operandNo,
-                                           otherHandle]() {
-          InFlightDiagnostic diag =
-              owner->emitOpError()
-              << "invalidated the handle to payload operations nested in the "
-                 "payload operation associated with its operand #"
-              << operandNo;
-          diag.attachNote(ancestorLoc) << "ancestor op";
-          diag.attachNote(opLoc) << "nested op";
-          diag.attachNote(otherHandle.getLoc()) << "other handle";
+                                           otherHandle](Location currentLoc) {
+          InFlightDiagnostic diag = emitError(currentLoc)
+                                    << "op uses a handle invalidated by a "
+                                       "previously executed transform op";
+          diag.attachNote(otherHandle.getLoc()) << "handle to invalidated ops";
+          diag.attachNote(owner->getLoc())
+              << "invalidated by this transform op that consumes its operand #"
+              << operandNo
+              << " and invalidates handles to payload ops nested in payload "
+                 "ops associated with the consumed handle";
+          diag.attachNote(ancestorLoc) << "ancestor payload op";
+          diag.attachNote(opLoc) << "nested payload op";
         };
       }
     }
   }
 }
 
 LogicalResult transform::TransformState::checkAndRecordHandleInvalidation(
     TransformOpInterface transform) {
   auto memoryEffectsIface =
       cast<MemoryEffectOpInterface>(transform.getOperation());
   SmallVector<MemoryEffects::EffectInstance> effects;
   memoryEffectsIface.getEffectsOnResource(
       transform::TransformMappingResource::get(), effects);
 
   for (OpOperand &target : transform->getOpOperands()) {
     // If the operand uses an invalidated handle, report it.
     auto it = invalidatedHandles.find(target.get());
     if (it != invalidatedHandles.end())
-      return it->getSecond()(), failure();
+      return it->getSecond()(transform->getLoc()), failure();
 
     // Invalidate handles pointing to the operations nested in the operation
     // associated with the handle consumed by this operation.
     auto consumesTarget = [&](const MemoryEffects::EffectInstance &effect) {
       return isa<MemoryEffects::Free>(effect.getEffect()) &&
              effect.getValue() == target.get();
     };
     if (llvm::any_of(effects, consumesTarget))

mlir/test/Dialect/Transform/expensive-checks.mlir

 // RUN: mlir-opt --test-transform-dialect-interpreter='enable-expensive-checks=1' --split-input-file --verify-diagnostics %s
 
-// expected-note @below {{ancestor op}}
+// expected-note @below {{ancestor payload op}}
 func.func @func() {
-  // expected-note @below {{nested op}}
+  // expected-note @below {{nested payload op}}
   return
 }
 
 transform.with_pdl_patterns {
 ^bb0(%arg0: !pdl.operation):
   pdl.pattern @return : benefit(1) {
     %0 = operands
     %1 = types
     %2 = operation "func.return"(%0 : !pdl.range<value>) -> (%1 : !pdl.range<type>)
     rewrite %2 with "transform.dialect"
   }
 
   sequence %arg0 failures(propagate) {
   ^bb1(%arg1: !pdl.operation):
-    // expected-note @below {{other handle}}
+    // expected-note @below {{handle to invalidated ops}}
     %0 = pdl_match @return in %arg1
     %1 = get_closest_isolated_parent %0
-    // expected-error @below {{invalidated the handle to payload operations nested in the payload operation associated with its operand #0}}
+    // expected-note @below {{invalidated by this transform op that consumes its operand #0}}
     test_consume_operand %1
+    // expected-error @below {{op uses a handle invalidated by a previously executed transform op}}
     test_print_remark_at_operand %0, "remark"
   }
 }
 
 // -----
 
 func.func @func1() {
   // expected-note @below {{repeated target op}}