[mlir] optionally allow repeated handles in transform dialect

[ftynse]

Summary

Some operations may be able to deal with handles pointing to the same
operation when the handle is consumed. For example, merge handles with
deduplication doesn't actually destroy payload operations and is
specifically intended to remove the situation with duplicates. Add a
method to the transform interface to allow ops to declare they can
support repeated handles.

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

InterfaceMethod<
         or failed, and, if it failed, whether the failure is recoverable or not.
       }],
       /*returnType=*/"::mlir::DiagnosedSilenceableFailure",
       /*name=*/"apply",
       /*arguments=*/(ins
           "::mlir::transform::TransformResults &":$transformResults,
           "::mlir::transform::TransformState &":$state
     )>,
+    InterfaceMethod<
+      /*desc=*/[{
+        Indicates whether the op instance allows its handle operands to be
+        associated with the same payload operations.
+      }],
+      /*returnType=*/"bool",
+      /*name=*/"allowsRepeatedHandleOperands",
+      /*arguments=*/(ins),
+      /*methodBody=*/"",
+      /*defaultImplementation=*/[{
+        return false;
+      }]
+    >,
   ];
 
   let extraSharedClassDeclaration = [{
     /// Creates the silenceable failure object with a diagnostic located at the
     /// current operation. Silenceable failure must be suppressed or reported
     /// explicitly at some later time.
     DiagnosedSilenceableFailure
     emitSilenceableError(const ::llvm::Twine &message = {}) {

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

def GetProducerOfOperand : TransformDialectOp<"get_producer_of_operand",
   let arguments = (ins TransformTypeInterface:$target,
                        I64Attr:$operand_number);
   let results = (outs TransformTypeInterface:$parent);
   let assemblyFormat = "$target `[` $operand_number `]` attr-dict `:` "
                        "functional-type(operands, results)";
 }
 
 def MergeHandlesOp : TransformDialectOp<"merge_handles",
-    [DeclareOpInterfaceMethods<TransformOpInterface>,
+    [DeclareOpInterfaceMethods<TransformOpInterface, ["allowsRepeatedHandleOperands"]>,
      DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
      SameOperandsAndResultType]> {
   let summary = "Merges handles into one pointing to the union of payload ops";
   let description = [{
     Creates a new Transform IR handle value that points to the same Payload IR
     operations as the operand handles. The Payload IR operations are listed
     in the same order as they are in the operand handles, grouped by operand
     handle, e.g., all Payload IR operations associated with the first handle

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

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())
+    if (!transform.allowsRepeatedHandleOperands() &&
+        it != invalidatedHandles.end())
       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))
       recordHandleInvalidation(target);
   }
+
   return success();
 }
 
 DiagnosedSilenceableFailure
 transform::TransformState::applyTransform(TransformOpInterface transform) {
   LLVM_DEBUG(DBGS() << "applying: " << transform << "\n");
   auto printOnFailureRAII = llvm::make_scope_exit([this] {
     (void)this;

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

if (!getDeduplicate()) {
     return DiagnosedSilenceableFailure::success();
   }
 
   SetVector<Operation *> uniqued(operations.begin(), operations.end());
   results.set(getResult().cast<OpResult>(), uniqued.getArrayRef());
   return DiagnosedSilenceableFailure::success();
 }
 
+bool transform::MergeHandlesOp::allowsRepeatedHandleOperands() {
+  // Handles may be the same if deduplicating is enabled.
+  return getDeduplicate();
+}
+
 void transform::MergeHandlesOp::getEffects(
     SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
   consumesHandle(getHandles(), effects);
   producesHandle(getResult(), effects);
 
   // There are no effects on the Payload IR as this is only a handle
   // manipulation.
 }

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

^bb0(%0: !pdl.operation):
     // Consuming two handles in the same operation is invalid if they point
     // to overlapping sets of payload IR ops.
     //
     // expected-error @below {{op uses a handle invalidated by a previously executed transform op}}
     // expected-note @below {{invalidated by this transform op that consumes its operand #0 and invalidates handles}}
     transform.test_consume_operand %1, %2
   }
 }
+
+// -----
+
+// Deduplication attribute allows "merge_handles" to take repeated operands.
+
+module {
+
+  transform.sequence failures(propagate) {
+  ^bb0(%0: !pdl.operation):
+    %1 = transform.test_copy_payload %0
+    %2 = transform.test_copy_payload %0
+    transform.merge_handles %1, %2 { deduplicate } : !pdl.operation
+  }
+}