[mlir][transform][NFC] Extract getConsumedHandleOpOperands helper function

[springerm]

Summary

This function is extracted from TransformState::applyTransform.

[ftynse]

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

1585

I see this was already there before this refactoring, but it doesn't look like this needs to be a set. OpOperands are unique by construction, so let's use a vector instead. It will also remove the implementation-defined behavior in traversals.

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

 /// the operands to `block`'s terminator.
 void forwardTerminatorOperands(Block *block, transform::TransformState &state,
                                transform::TransformResults &results);
 
 /// Make a dummy transform state for testing purposes. This MUST NOT be used
 /// outside of test cases.
 TransformState makeTransformStateForTesting(Region *region,
                                             Operation *payloadRoot);
+
+/// Returns all operands that are handles and being consumed by the given op.
+SmallVector<OpOperand *>
+getConsumedHandleOpOperands(transform::TransformOpInterface transformOp);
 } // namespace detail
 
 /// Options controlling the application of transform operations by the
 /// TransformState.
 class TransformOptions {
 public:
   TransformOptions() = default;
   TransformOptions(const TransformOptions &) = default;

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

let extraSharedClassDeclaration = [{
     /// to properly apply to a target.
     DiagnosedSilenceableFailure emitDefaultSilenceableFailure(
         Operation *target) {
       DiagnosedSilenceableFailure diag = emitSilenceableFailure($_op->getLoc());
       diag << $_op->getName() << " failed to apply";
       diag.attachNote(target->getLoc()) << "when applied to this op";
       return diag;
     }
+
+    /// Returns all operands that are handles and being consumed by this op.
+    ::llvm::SmallVector<OpOperand *> getConsumedHandleOpOperands() {
+      return ::mlir::transform::detail::getConsumedHandleOpOperands($_op);
+    }
   }];
 
   let verify = [{
     return ::mlir::transform::detail::verifyTransformOpInterface($_op);
   }];
 }
 
 class TransformTypeInterfaceBase<string cppClass, string cppObjectType>

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

for (Mappings &mapping : llvm::make_second_range(mappings)) {
           }
         }
       }
     }
 #endif // LLVM_ENABLE_ABI_BREAKING_CHECKS
   }
 
   // Find which operands are consumed.
-  DenseSet<unsigned> consumedOperands;
-  auto memEffectInterface =
-      cast<MemoryEffectOpInterface>(transform.getOperation());
-  SmallVector<MemoryEffects::EffectInstance, 2> effects;
-  for (OpOperand &target : transform->getOpOperands()) {
-    effects.clear();
-    memEffectInterface.getEffectsOnValue(target.get(), effects);
-    if (llvm::any_of(effects, [](const MemoryEffects::EffectInstance &effect) {
-          return isa<transform::TransformMappingResource>(
-                     effect.getResource()) &&
-                 isa<MemoryEffects::Free>(effect.getEffect());
-        })) {
-      consumedOperands.insert(target.getOperandNumber());
-    }
-  }
+  SmallVector<OpOperand *> consumedOperands =
+      transform.getConsumedHandleOpOperands();
 
   // 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;
 #if LLVM_ENABLE_ABI_BREAKING_CHECKS
   DenseSet<Operation *> consumedPayloadOps;
 #endif // LLVM_ENABLE_ABI_BREAKING_CHECKS
-  for (unsigned index : consumedOperands) {
-    Value operand = transform->getOperand(index);
+  for (OpOperand *opOperand : consumedOperands) {
+    Value operand = opOperand->get();
     if (llvm::isa<TransformHandleTypeInterface>(operand.getType())) {
       for (Operation *payloadOp : getPayloadOps(operand)) {
         llvm::append_range(origOpFlatResults, payloadOp->getResults());
 #if LLVM_ENABLE_ABI_BREAKING_CHECKS
         if (options.getExpensiveChecksEnabled()) {
           // Store all consumed payload ops (and their nested ops) in a set for
           // extra error checking.
           payloadOp->walk(
if (llvm::isa<TransformValueHandleTypeInterface>(operand.getType())) {
             llvm::map_range(*llvm::cast<BlockArgument>(payloadValue).getOwner(),
                             [](Operation &op) { return &op; }));
       }
       continue;
     }
     DiagnosedDefiniteFailure diag =
         emitDefiniteFailure(transform->getLoc())
         << "unexpectedly consumed a value that is not a handle as operand #"
-        << index;
+        << opOperand->getOperandNumber();
     diag.attachNote(operand.getLoc())
         << "value defined here with type " << operand.getType();
     return diag;
   }
 
   // Compute the result but do not short-circuit the silenceable failure case as
   // we still want the handles to propagate properly so the "suppress" mode can
   // proceed on a best effort basis.
   transform::TransformResults results(transform->getNumResults());
   DiagnosedSilenceableFailure result(transform.apply(results, *this));
   compactOpHandles();
   if (result.isDefiniteFailure())
     return result;
 
   // If a silenceable failure was produced, some results may be unset, set them
   // to empty lists.
   if (result.isSilenceableFailure())
     results.setRemainingToEmpty(transform);
 
   // Remove the mapping for the operand if it is consumed by the operation. This
   // allows us to catch use-after-free with assertions later on.
-  for (unsigned index : consumedOperands) {
-    Value operand = transform->getOperand(index);
+  for (OpOperand *opOperand : consumedOperands) {
+    Value operand = opOperand->get();
     if (llvm::isa<TransformHandleTypeInterface>(operand.getType())) {
       forgetMapping(operand, origOpFlatResults);
     } else if (llvm::isa<TransformValueHandleTypeInterface>(
                    operand.getType())) {
       forgetValueMapping(operand, origAssociatedOps);
     }
   }
 
for (Operation &nested : block) {
     }
   }
 }
 
 //===----------------------------------------------------------------------===//
 // Utilities for TransformOpInterface.
 //===----------------------------------------------------------------------===//
 
+SmallVector<OpOperand *> transform::detail::getConsumedHandleOpOperands(
+    TransformOpInterface transformOp) {
+  SmallVector<OpOperand *> consumedOperands;
+  consumedOperands.reserve(transformOp->getNumOperands());
+  auto memEffectInterface =
+      cast<MemoryEffectOpInterface>(transformOp.getOperation());
+  SmallVector<MemoryEffects::EffectInstance, 2> effects;
+  for (OpOperand &target : transformOp->getOpOperands()) {
+    effects.clear();
+    memEffectInterface.getEffectsOnValue(target.get(), effects);
+    if (llvm::any_of(effects, [](const MemoryEffects::EffectInstance &effect) {
+          return isa<transform::TransformMappingResource>(
+                     effect.getResource()) &&
+                 isa<MemoryEffects::Free>(effect.getEffect());
+        })) {
+      consumedOperands.push_back(&target);
+    }
+  }
+  return consumedOperands;
+}
+
 LogicalResult transform::detail::verifyTransformOpInterface(Operation *op) {
   auto iface = cast<MemoryEffectOpInterface>(op);
   SmallVector<MemoryEffects::EffectInstance> effects;
   iface.getEffects(effects);
 
   auto effectsOn = [&](Value value) {
     return llvm::make_filter_range(
         effects, [value](const MemoryEffects::EffectInstance &instance) {