[mlir] teach expensive-checks transform mode about empty handle

[ftynse]

Summary

The transform dialect interpreter features the expensive-checks mode
that acts as an embedded sanitizer to track use-after-consume of
transform handles. Its logic is based on the relations between payload
operations, which made it silently ignore empty handles that are
consumed. Also catch and report this case because the remaining code may
hit an assertion on attempting to access a consumed handle (that is
removed from the mapping).

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

if (invalidatedHandles.count(otherHandle))
     return;
 
   FULL_LDBG("--recordOpHandleInvalidationOne\n");
   DEBUG_WITH_TYPE(
       DEBUG_TYPE_FULL,
       llvm::interleaveComma(potentialAncestors, DBGS() << "--ancestors: ",
                             [](Operation *op) { llvm::dbgs() << *op; });
       llvm::dbgs() << "\n");
+
+  Operation *owner = consumingHandle.getOwner();
+  unsigned operandNo = consumingHandle.getOperandNumber();
   for (Operation *ancestor : potentialAncestors) {
     // clang-format off
     DEBUG_WITH_TYPE(DEBUG_TYPE_FULL, 
       { (DBGS() << "----handle one ancestor: " << *ancestor << "\n"); });
     DEBUG_WITH_TYPE(DEBUG_TYPE_FULL, 
       { (DBGS() << "----of payload with name: " 
                 << payloadOp->getName().getIdentifier() << "\n"); });
     DEBUG_WITH_TYPE(DEBUG_TYPE_FULL,
       { (DBGS() << "----of payload: " << *payloadOp << "\n"); });
     // clang-format on
     if (!ancestor->isAncestor(payloadOp))
       continue;
 
     // Make sure the error-reporting lambda doesn't capture anything
     // 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 = payloadOp->getLoc();
-    Operation *owner = consumingHandle.getOwner();
-    unsigned operandNo = consumingHandle.getOperandNumber();
     std::optional<Location> throughValueLoc =
         throughValue ? std::make_optional(throughValue.getLoc()) : std::nullopt;
     invalidatedHandles[otherHandle] = [ancestorLoc, opLoc, owner, operandNo,
                                        otherHandle,
                                        throughValueLoc](Location currentLoc) {
       InFlightDiagnostic diag = emitError(currentLoc)
                                 << "op uses a handle invalidated by a "
                                    "previously executed transform op";
invalidatedHandles[valueHandle] =
           diag.attachNote(valueLoc) << "payload value";
         };
   }
 }
 
 void transform::TransformState::recordOpHandleInvalidation(
     OpOperand &handle, ArrayRef<Operation *> potentialAncestors,
     Value throughValue) {
+
+  if (potentialAncestors.empty()) {
+    DEBUG_WITH_TYPE(DEBUG_TYPE_FULL, {
+      (DBGS() << "----recording invalidation for empty handle: " << handle.get()
+              << "\n");
+    });
+
+    Operation *owner = handle.getOwner();
+    unsigned operandNo = handle.getOperandNumber();
+    invalidatedHandles[handle.get()] = [owner, operandNo](Location currentLoc) {
+      InFlightDiagnostic diag = emitError(currentLoc)
+                                << "op uses a handle associated with empty "
+                                   "payload and invalidated by a "
+                                   "previously executed transform op";
+      diag.attachNote(owner->getLoc())
+          << "invalidated by this transform op that consumes its operand #"
+          << operandNo;
+    };
+    return;
+  }
+
   // Iterate over the mapping and invalidate aliasing handles. This is quite
   // expensive and only necessary for error reporting in case of transform
   // dialect misuse with dangling handles. Iteration over the handles is based
   // on the assumption that the number of handles is significantly less than the
   // number of IR objects (operations and values). Alternatively, we could walk
   // the IR nested in each payload op associated with the given handle and look
   // for handles associated with each operation and value.
   for (const Mappings &mapping : llvm::make_second_range(mappings)) {

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

 transform.sequence failures(propagate) {
 ^bb1(%arg0: !transform.any_op):
   %1 = transform.structured.match ops{["test.match_anchor_1"]} in %arg0 : (!transform.any_op) -> !transform.any_op
   %2 = transform.structured.match ops{["test.match_anchor_2"]} in %arg0 : (!transform.any_op) -> !transform.any_op
   %3 = test_produce_value_handle_to_argument_of_parent_block %1, 0 : (!transform.any_op) -> !transform.any_value
   test_consume_operand %3 : !transform.any_value
   test_consume_operand %2 : !transform.any_op
 }
+
+// -----
+
+transform.sequence failures(propagate) {
+^bb0(%arg0: !transform.any_op):
+  %0 = transform.test_produce_empty_payload : !transform.any_op
+  // expected-note @below {{invalidated by this transform op that consumes its operand #0}}
+  transform.test_consume_operand %0 : !transform.any_op
+  // expected-error @below {{uses a handle associated with empty payload and invalidated by a previously executed transform op}}
+  transform.test_print_remark_at_operand %0, "remark" : !transform.any_op
+}

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

 
 DiagnosedSilenceableFailure mlir::test::TestProduceNullPayloadOp::apply(
     transform::TransformResults &results, transform::TransformState &state) {
   SmallVector<Operation *, 1> null({nullptr});
   results.set(llvm::cast<OpResult>(getOut()), null);
   return DiagnosedSilenceableFailure::success();
 }
 
+DiagnosedSilenceableFailure mlir::test::TestProduceEmptyPayloadOp::apply(
+    transform::TransformResults &results, transform::TransformState &state) {
+  results.set(cast<OpResult>(getOut()), {});
+  return DiagnosedSilenceableFailure::success();
+}
+
 void mlir::test::TestProduceNullParamOp::getEffects(
     SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
   transform::producesHandle(getOut(), effects);
 }
 
 DiagnosedSilenceableFailure
 mlir::test::TestProduceNullParamOp::apply(transform::TransformResults &results,
                                           transform::TransformState &state) {

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

def TestProduceNullPayloadOp
   : Op<Transform_Dialect, "test_produce_null_payload",
       [DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
        DeclareOpInterfaceMethods<TransformOpInterface>]> {
   let results = (outs TransformHandleTypeInterface:$out);
   let assemblyFormat = "attr-dict `:` type($out)";
   let cppNamespace = "::mlir::test";
 }
 
+def TestProduceEmptyPayloadOp
+  : Op<Transform_Dialect, "test_produce_empty_payload",
+      [DeclareOpInterfaceMethods<TransformOpInterface>,
+       MemoryEffectsOpInterface, FunctionalStyleTransformOpTrait]> {
+  let results = (outs TransformHandleTypeInterface:$out);
+  let assemblyFormat = "attr-dict `:` type($out)";
+  let cppNamespace = "::mlir::test";
+}
+
 def TestProduceNullParamOp
   : Op<Transform_Dialect, "test_produce_null_param",
       [DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
        DeclareOpInterfaceMethods<TransformOpInterface>]> {
   let results = (outs TransformParamTypeInterface:$out);
   let assemblyFormat = "attr-dict `:` type($out)";
   let cppNamespace = "::mlir::test";
 }