[mlir] fix side effects for transform.AlternativesOp

[ftynse]

Summary

It should have an "Allocate" effect on entry block arguments of all
regions in addition to consuming the operand.

Also relax the assertion in transform-dialect-check-uses until we can
properly support region-based control flow.

Fixes #60075.

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

 include "mlir/Dialect/Transform/IR/TransformEffects.td"
 include "mlir/Dialect/Transform/IR/TransformInterfaces.td"
 
 def AlternativesOp : TransformDialectOp<"alternatives",
     [DeclareOpInterfaceMethods<RegionBranchOpInterface,
         ["getSuccessorEntryOperands", "getSuccessorRegions",
          "getRegionInvocationBounds"]>,
      DeclareOpInterfaceMethods<TransformOpInterface>,
-     FunctionalStyleTransformOpTrait, MemoryEffectsOpInterface,
+     DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
      IsolatedFromAbove, PossibleTopLevelTransformOpTrait,
      SingleBlockImplicitTerminator<"::mlir::transform::YieldOp">]> {
   let summary = "Attempts sequences of transforms until one succeeds";
   let description = [{
     This op may have an arbitrary number of regions, each of which represents a
     sequence of transform operations to be applied to the same payload IR. The
     regions are visited in order of appearance, and transforms in them are
     applied in their respective order of appearance. If one of these transforms

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

if (!failed) {
       }
       forwardTerminatorOperands(&reg.front(), state, results);
       return DiagnosedSilenceableFailure::success();
     }
   }
   return emitSilenceableError() << "all alternatives failed";
 }
 
+void transform::AlternativesOp::getEffects(
+    SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
+  consumesHandle(getOperands(), effects);
+  producesHandle(getResults(), effects);
+  for (Region *region : getRegions()) {
+    if (!region->empty())
+      producesHandle(region->front().getArguments(), effects);
+  }
+  modifiesPayload(effects);
+}
+
 LogicalResult transform::AlternativesOp::verify() {
   for (Region &alternative : getAlternatives()) {
     Block &block = alternative.front();
     Operation *terminator = block.getTerminator();
     if (terminator->getOperands().getTypes() != getResults().getTypes()) {
       InFlightDiagnostic diag = emitOpError()
                                 << "expects terminator operands to have the "
                                    "same type as results of the operation";

mlir/lib/Dialect/Transform/Transforms/CheckUses.cpp

public:
   /// the given operand. For the purposes of this analysis, the value is
   /// considered to be allocated at its definition point and never re-allocated.
   PotentialDeleters isUseLive(OpOperand &operand) {
     const llvm::SmallPtrSet<Operation *, 2> &deleters = freedBy[operand.get()];
     if (deleters.empty())
       return live();
 
 #ifndef NDEBUG
-    // Check that the definition point actually allcoates the value.
+    // Check that the definition point actually allocates the value. If the
+    // definition is a block argument, it may be just forwarding the operand of
+    // the parent op without doing a new allocation, allow that. We currently
+    // don't have the capability to analyze region-based control flow here.
+    //
+    // TODO: when this ported to the dataflow analysis infra, we should have
+    // proper support for region-based control flow.
     Operation *valueSource =
         operand.get().isa<OpResult>()
             ? operand.get().getDefiningOp()
             : operand.get().getParentBlock()->getParentOp();
     auto iface = cast<MemoryEffectOpInterface>(valueSource);
     SmallVector<MemoryEffects::EffectInstance> instances;
     iface.getEffectsOnResource(transform::TransformMappingResource::get(),
                                instances);
-    assert(hasEffect<MemoryEffects::Allocate>(instances, operand.get()) &&
+    assert((operand.get().isa<BlockArgument>() ||
+            hasEffect<MemoryEffects::Allocate>(instances, operand.get())) &&
            "expected the op defining the value to have an allocation effect "
            "on it");
 #endif
 
     // Collect ancestors of the use operation.
     Block *defBlock = operand.get().getParentBlock();
     SmallVector<Operation *> ancestors;
     Operation *ancestor = operand.getOwner();

mlir/test/Dialect/Transform/check-use-after-free.mlir

func.func @use_after_free_cycle() {
   ^bb2:
     "transform.test_branching_transform_op_terminator"()[^bb1] : () -> ()
   ^bb3:
     "transform.test_branching_transform_op_terminator"() : () -> ()
   }
   return
 }
 
+// -----
+
+// This should not crash.
+
+transform.sequence failures(propagate) {
+^bb0(%arg0: !pdl.operation):
+  alternatives %arg0 : !pdl.operation {
+  ^bb0(%arg1: !pdl.operation):
+  }
+}