[mlir][transform] Better debugging facilites for invalid API usage

[springerm]

Summary

This revision adds additional "expensive-checks" checks to the transform dialect that detect the most common cases of:

• Missing consumesHandle side effects on transform ops.
• Patterns that remove operations but do not notify the transform dialect.

In essence, these additional checks are looking for dangling pointers to erased payload ops in the transform dialect state and crash the program execution (by dereferencing free'd memory) or triggering an assertion failure. It is recommended to run these extra checks with ASAN. Otherwise, certain failures may not be detected. The ASAN error message can also be used to find the faulty transform op/pattern.

This change also fixes a few faulty transform ops.

Depends On: D147446

[springerm]

update

[nicolasvasilache]

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

689

spew

[ftynse]

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

568

This should be LLVM_ENABLE_ABI_BREAKING_CHECKS because it changes the size of the object.

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

745

Nit: if only the second part is needed, use llvm::make_second_range. This will also allow auto to be expanded for better readabiltiy.

746

Avoid shadowing iterator names, this is terrible for readability and may backfire if the loop is no longer a range-for.

783

Can we have something similar for values? Consuming an op also consumes values defined by this op (op results or block arguments).

855

Nit: can we factor this loop out into a function so the parent function doesn't get excessively long?

857–860

Same as above (make_second_range and not reusing it for two loops).

873–878

This looks similar to forgetMapping, can we refactor/reuse that somehow. Specifically, this creates a situation where a handle can be associated with zero payload operations, which usually results in the handle being dropped completely, but not here. Some refactoring is likely necessary because we cannot drop the handle from the direct mapping without invalidating the iterator used in the outer loop.

878

On a second thought, it looks reasonable to keep those handles so we report a nice diagnostic about using an invalidated handle instead of asserting on "unknown handle". (And we should update the other place not to assert.) However, let's have some more explanation as to why this doesn't interfere with the other expensive check: handle invalidation happens before this and keeps track of invalidated handles with pregenerated error messages, so we don't need the association to still be there when the invalidated handle is accessed.

887–891

I'm hesitating about making these definite errors instead. This can happen in "syntactically valid but semantically invalid" transform IR that we get from the user and have no other way of checking for validity, so not crashing sounds preferable. However, we will still crash in the dangling-pointer case and there is no workaround for that.

[springerm]

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

783

I'm going to do that as part of a second revision that adds the same "expensive check" for OpResults. I think we need a second cachedNames.

855

I'd prefer to keep it here because consumedPayloadOps is also defined in this function. Otherwise, I have to add another field to TransformState (or pass additional parameters around).

878

Using forgetHandle actually cleans up the code quite a bit.

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

private:
   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
+
+#ifndef NDEBUG
+  /// This cache stores operation names for operations that are tracked in the
+  /// transform dialect state. It is used to detect missing memory side effects
+  /// and op tracking.
+  ///
+  /// All tracked ops are added to this cache before a transform op is applied.
+  /// After the application of the transform op, the names of all tracked ops
+  /// are compared with the names in the cache. If there is a mismatch (or a
+  /// crash), op tracking is missing somewhere. This is typically a missing
+  /// "consumesHandle" side effect or a pattern that removes an op without
+  /// notifying a TrackingListener.
+  DenseMap<Operation *, OperationName> cachedNames;
+#endif // NDEBUG
 };
 
 /// Local mapping between values defined by a specific op implementing the
 /// TransformOpInterface and the payload IR ops they correspond to.
 class TransformResults {
   friend class TransformState;
 
 public:

mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp

 //===---------------------------------------------------------------------===//
 
 DiagnosedSilenceableFailure
 transform::HoistPadOp::applyToOne(tensor::PadOp target,
                                   transform::ApplyToEachResultList &results,
                                   transform::TransformState &state) {
   tensor::PadOp hoistedPadOp;
   SmallVector<GenericOp> transposeOps;
-  IRRewriter rewriter(target->getContext());
+  TrackingListener listener(state);
+  IRRewriter rewriter(target->getContext(), &listener);
   FailureOr<Value> result =
       hoistPaddingOnTensors(rewriter, target, getNumLoops(), getTranspose(),
                             hoistedPadOp, transposeOps);
   if (succeeded(result)) {
     // We need to perform our own replacement here because this API is still
     // used in patterns that "pad and hoist", for which the replacement values
     // need to be different.
     // TODO: clean this up and stop "pad and hoist" behavior more globally now

mlir/lib/Dialect/Linalg/Transforms/HoistPadding.cpp

 ///   5. The backward slice from the pad op to the scf::ForOp to hoist above is
 ///      empty.
 ///   6. The source tensor of pad op is not defined by an extract slice op.
 ///   7. The source tensor of the extract slice op is not defined outside of
 ///      the outermost enclosing scf::ForOp.
 ///   8. There is no enclosing scf::ForOp that indexes the padded data.
 /// Other cases succeed and will trigger hoisting of the pad op.
 struct HoistingAnalysis {
-  HoistingAnalysis(tensor::PadOp padOp, int numLoops);
+  HoistingAnalysis(tensor::PadOp padOp, int numLoops,
+                   OpBuilder::Listener *listener);
 
   bool isValid() { return valid; }
 
   /// Footprint of the packedTensor, computed from the packingLoops.
   SmallVector<Value> getPackedTensorSizes(RewriterBase &b, Location loc);
 
   /// The outermost loop, determined by `nLevels` above which `padOp` will
   /// be hoisted.
static void computeBackwardSlice(tensor::PadOp padOp,
   for (Value v : valuesDefinedAbove) {
     getBackwardSlice(v, &backwardSlice, filter, /*inclusive=*/true);
   }
   // Then, add the backward slice from padOp itself.
   getBackwardSlice(padOp.getOperation(), &backwardSlice, filter,
                    /*inclusive=*/true);
 }
 
-HoistingAnalysis::HoistingAnalysis(tensor::PadOp padOp, int numLoops) {
+HoistingAnalysis::HoistingAnalysis(tensor::PadOp padOp, int numLoops,
+                                   OpBuilder::Listener *listener) {
   valid = false;
 
   // Get at most `numLoops` of immediately enclosing loops.
   SmallVector<scf::ForOp> reverseEnclosingLoops;
   getAtMostNEnclosingLoops(padOp, numLoops, reverseEnclosingLoops);
   if (reverseEnclosingLoops.empty()) {
     LLVM_DEBUG(DBGS() << "--No immediately enclosing loop -> Skip\n");
     return;
HoistingAnalysis::HoistingAnalysis(tensor::PadOp padOp, int numLoops,
   sliceOp = padOp.getSource().getDefiningOp<tensor::ExtractSliceOp>();
   if (!sliceOp) {
     LLVM_DEBUG(DBGS() << "--Cannot find the extract slice op -> Skip\n");
     return;
   }
   // If the padded data is not yet available before entering the outermost
   // enclosing loop, try to apply hoisting on this outermost loop.
   // TODO: we may want finer-grained hoisting of only that particular `sliceOp`.
-  IRRewriter rewriter(outermostEnclosingForOp->getContext());
+  IRRewriter rewriter(outermostEnclosingForOp->getContext(), listener);
   if (!outermostEnclosingForOp.isDefinedOutsideOfLoop(sliceOp.getSource())) {
     outermostEnclosingForOp =
         hoistRedundantSubsetExtractInsert(rewriter, outermostEnclosingForOp);
   }
   if (!outermostEnclosingForOp.isDefinedOutsideOfLoop(sliceOp.getSource())) {
     LLVM_DEBUG(DBGS() << "--outermostEnclosingForOp:\n"
                       << outermostEnclosingForOp << "\n"
                       << "--sliceOp: " << sliceOp << "\n"

 }
 
 FailureOr<Value> mlir::linalg::hoistPaddingOnTensors(
     RewriterBase &rewriter, tensor::PadOp opToHoist, int64_t numLoops,
     ArrayRef<int64_t> transposeVector, tensor::PadOp &hoistedOp,
     SmallVectorImpl<GenericOp> &transposeOps) {
   LLVM_DEBUG(DBGS() << "\n"; DBGS() << " Try to hoist " << *(opToHoist) << "\n";
              DBGS() << " by " << numLoops << " loops\n");
-  HoistingAnalysis analysis(opToHoist, numLoops);
+  HoistingAnalysis analysis(opToHoist, numLoops, rewriter.getListener());
   if (!analysis.isValid()) {
     LLVM_DEBUG(DBGS() << "--Analysis failed -> Skip\n");
     return failure();
   }
 
   // Update actual number of loops, which may be smaller.
   int nPackedLoops = analysis.packingLoops.size();
   LLVM_DEBUG(DBGS() << "\n";

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

transform::TransformState::replacePayloadOp(Operation *op,
   }
 
 #ifndef NDEBUG
   for (Value opResult : op->getResults()) {
     SmallVector<Value> valueHandles;
     (void)getHandlesForPayloadValue(opResult, valueHandles);
     assert(valueHandles.empty() && "expected no mapping to old results");
   }
+
+  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 // NDEBUG
 
   // TODO: consider invalidating the handles to nested objects here.
 
   // If replacing with null, that is erasing the mapping, drop the mapping
   // between the handles and the IR objects and return.
   if (!replacement) {
     for (Value handle : opHandles) {
       Mappings &mappings = getMapping(handle);
       dropMappingEntry(mappings.direct, handle, op);
     }
     return success();
   }
 
+#ifndef NDEBUG
+  if (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 // NDEBUG
+
   // Otherwise, replace the pointed-to object of all handles while preserving
   // their relative order. First, replace the mapped operation if present.
   for (Value handle : opHandles) {
     Mappings &mappings = getMapping(handle);
     auto it = mappings.direct.find(handle);
     if (it == mappings.direct.end())
       continue;
 
if (!seen.insert(p).second) {
       return diag;
     }
   }
   return DiagnosedSilenceableFailure::success();
 }
 
 DiagnosedSilenceableFailure
 transform::TransformState::applyTransform(TransformOpInterface transform) {
+  llvm::errs() << "apply transform: " << transform << "\n";
+
   LLVM_DEBUG(DBGS() << "\n"; DBGS() << "applying: " << transform << "\n");
   LLVM_DEBUG(DBGS() << "On top-level payload:\n" << *getTopLevel(););
   auto printOnFailureRAII = llvm::make_scope_exit([this] {
     (void)this;
     LLVM_DEBUG(DBGS() << "Failing Top-level payload:\n"; getTopLevel()->print(
         llvm::dbgs(), mlir::OpPrintingFlags().printGenericOpForm()););
   });
+
   if (options.getExpensiveChecksEnabled()) {
     FULL_LDBG("ExpensiveChecksEnabled\n");
     if (failed(checkAndRecordHandleInvalidation(transform)))
       return DiagnosedSilenceableFailure::definiteFailure();
 
     for (OpOperand &operand : transform->getOpOperands()) {
       DEBUG_WITH_TYPE(DEBUG_TYPE_FULL, {
         (DBGS() << "iterate on handle: " << operand.get() << "\n");
for (OpOperand &operand : transform->getOpOperands()) {
         if (!check.succeeded()) {
           FULL_LDBG("----FAILED\n");
           return check;
         }
       } else {
         FULL_LDBG("--not a TransformHandle -> SKIP AND DROP ON THE FLOOR\n");
       }
     }
+
+#ifndef NDEBUG
+    // Cache Operation* -> OperationName mappings. These will be checked after
+    // the transform has been applied to detect incorrect memory side effects
+    // and missing op tracking.
+    for (auto &it : mappings) {
+      Mappings &mapping = it.second;
+      for (auto &it : mapping.reverse) {
+        Operation *op = it.first;
+        auto insertion = cachedNames.insert({op, op->getName()});
+        if (!insertion.second) {
+          assert(insertion.first->second == op->getName() &&
+                 "operation is already cached with a different name");
+        }
+      }
+    }
+#endif // NDEBUG
   }
 
   // Find which operands are consumed.
   DenseSet<unsigned> consumedOperands;
   auto memEffectInterface =
       cast<MemoryEffectOpInterface>(transform.getOperation());
   SmallVector<MemoryEffects::EffectInstance, 2> effects;
   for (OpOperand &target : transform->getOpOperands()) {
#endif // NDEBUG
 
   // Remember the results of the payload ops associated with the consumed
   // op handles or the ops defining the value handles so we can drop the
   // association with them later. This must happen here because the
   // transformation may destroy or mutate them so we cannot traverse the payload
   // IR after that.
   SmallVector<Value> origOpFlatResults;
   SmallVector<Operation *> origAssociatedOps;
+  DenseSet<Operation *> consumedPayloadOps;
   for (unsigned index : consumedOperands) {
     Value operand = transform->getOperand(index);
     if (operand.getType().isa<TransformHandleTypeInterface>()) {
-      for (Operation *payloadOp : getPayloadOps(operand))
+      for (Operation *payloadOp : getPayloadOps(operand)) {
         llvm::append_range(origOpFlatResults, payloadOp->getResults());
+        if (options.getExpensiveChecksEnabled()) {
+          // Store all consumed payload ops (and their nested ops) in a set for
+          // extra error checking.
+          payloadOp->walk(
+              [&](Operation *op) { consumedPayloadOps.insert(op); });
+        }
+      }
       continue;
     }
     if (operand.getType().isa<TransformValueHandleTypeInterface>()) {
       for (Value payloadValue : getPayloadValues(operand)) {
         if (payloadValue.isa<OpResult>()) {
           origAssociatedOps.push_back(payloadValue.getDefiningOp());
           continue;
         }
for (unsigned index : consumedOperands) {
     Value operand = transform->getOperand(index);
     if (operand.getType().isa<TransformHandleTypeInterface>()) {
       forgetMapping(operand, origOpFlatResults);
     } else if (operand.getType().isa<TransformValueHandleTypeInterface>()) {
       forgetValueMapping(operand, origAssociatedOps);
     }
   }
 
+#ifndef NDEBUG
+  if (options.getExpensiveChecksEnabled()) {
+    // Check cached operation names.
+    DominanceInfo domInfo;
+    for (auto &it : mappings) {
+      Mappings &mapping = it.second;
+      for (auto &it : mapping.reverse) {
+        Operation *op = it.first;
+        if (consumedPayloadOps.contains(op)) {
+          // This payload op was consumed but it is still mapped to one or
+          // multiple handles. Erase the op from all mappings, so that there are
+          // no dangling pointers in the transform dialect state.
+          for (Value handle : it.second) {
+            auto it = llvm::find(mapping.direct[handle], op);
+            assert(it != mapping.direct[handle].end() &&
+                   "inconsistent mapping state");
+            mapping.direct[handle].erase(it);
+          }
+          mapping.reverse.erase(op);
+          cachedNames.erase(op);
+          continue;
+        }
+        // Make sure that the name of the op has not changed. If it has changed,
+        // the op was removed and a new op was allocated at the same memory
+        // location. This means that we are missing op tracking somewhere.
+        auto cacheIt = cachedNames.find(op);
+        assert(cacheIt != cachedNames.end() && "operation not in cache");
+        assert(cacheIt->second == op->getName() && "operation name mismatch");
+      }
+    }
+  }
+#endif // NDEBUG
+
   for (OpResult result : transform->getResults()) {
     assert(result.getDefiningOp() == transform.getOperation() &&
            "payload IR association for a value other than the result of the "
            "current transform op");
     if (result.getType().isa<TransformParamTypeInterface>()) {
       assert(results.isParam(result.getResultNumber()) &&
              "expected parameters for the parameter-typed result");
       if (failed(
               setParams(result, results.getParams(result.getResultNumber())))) {
         return DiagnosedSilenceableFailure::definiteFailure();
       }
     } else if (result.getType().isa<TransformValueHandleTypeInterface>()) {
       assert(results.isValue(result.getResultNumber()) &&

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

for (Region &reg : getAlternatives()) {
     }
 
     // If all operations in the given alternative succeeded, no need to consider
     // the rest. Replace the original scoping operation with the clone on which
     // the transformations were performed.
     if (!failed) {
       // We will be using the clones, so cancel their scheduled deletion.
       deleteClones.release();
-      IRRewriter rewriter(getContext());
+      TrackingListener listener(state);
+      IRRewriter rewriter(getContext(), &listener);
       for (const auto &kvp : llvm::zip(originals, clones)) {
         Operation *original = std::get<0>(kvp);
         Operation *clone = std::get<1>(kvp);
         original->getBlock()->getOperations().insert(original->getIterator(),
                                                      clone);
         rewriter.replaceOp(original, clone->getResults());
       }
       forwardTerminatorOperands(&reg.front(), state, results);

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

DiagnosedSilenceableFailure mlir::test::TestEmitRemarkAndEraseOperandOp::apply(
   for (Operation *op : state.getPayloadOps(getTarget()))
     op->erase();
 
   if (getFailAfterErase())
     return emitSilenceableError() << "silenceable error";
   return DiagnosedSilenceableFailure::success();
 }
 
+void mlir::test::TestEmitRemarkAndEraseOperandOp::getEffects(
+    SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
+  transform::consumesHandle(getTarget(), effects);
+  transform::modifiesPayload(effects);
+}
+
 DiagnosedSilenceableFailure mlir::test::TestWrongNumberOfResultsOp::applyToOne(
     Operation *target, transform::ApplyToEachResultList &results,
     transform::TransformState &state) {
   OperationState opState(target->getLoc(), "foo");
   results.push_back(OpBuilder(target).create(opState));
   return DiagnosedSilenceableFailure::success();
 }
 

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

: Op<Transform_Dialect, "test_branching_transform_op_terminator",
         DeclareOpInterfaceMethods<MemoryEffectsOpInterface>]> {
   let successors = (successor VariadicSuccessor<AnySuccessor>:$succ);
   let cppNamespace = "::mlir::test";
 }
 
 def TestEmitRemarkAndEraseOperandOp
   : Op<Transform_Dialect, "test_emit_remark_and_erase_operand",
     [DeclareOpInterfaceMethods<TransformOpInterface>,
-     MemoryEffectsOpInterface, FunctionalStyleTransformOpTrait]> {
+     DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
+     FunctionalStyleTransformOpTrait]> {
   let arguments = (ins PDL_Operation:$target, StrAttr:$remark,
                    UnitAttr:$fail_after_erase);
   let assemblyFormat = "$target `,` $remark attr-dict";
   let cppNamespace = "::mlir::test";
 }
 
 def TestWrongNumberOfResultsOp
   : Op<Transform_Dialect, "test_wrong_number_of_results",