[Mlir-commits] [mlir] d46afee - [mlir] fix side effects for transform.AlternativesOp

[Alex Zinenko]

Author: Alex Zinenko
Date: 2023-01-23T15:17:09Z
New Revision: d46afeef739a6efa71a3a78c52e7ed0b89a51725

URL: https://github.com/llvm/llvm-project/commit/d46afeef739a6efa71a3a78c52e7ed0b89a51725
DIFF: https://github.com/llvm/llvm-project/commit/d46afeef739a6efa71a3a78c52e7ed0b89a51725.diff

LOG: [mlir] fix side effects for transform.AlternativesOp

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.

Reviewed By: springerm

Differential Revision: https://reviews.llvm.org/D142200

Added: 
    

Modified: 
    mlir/include/mlir/Dialect/Transform/IR/TransformOps.td
    mlir/lib/Dialect/Transform/IR/TransformOps.cpp
    mlir/lib/Dialect/Transform/Transforms/CheckUses.cpp
    mlir/test/Dialect/Transform/check-use-after-free.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Dialect/Transform/IR/TransformOps.td b/mlir/include/mlir/Dialect/Transform/IR/TransformOps.td
index 699929fa17f0a..2dc93c1fb7ef3 100644
--- a/mlir/include/mlir/Dialect/Transform/IR/TransformOps.td
+++ b/mlir/include/mlir/Dialect/Transform/IR/TransformOps.td
@@ -24,7 +24,7 @@ def AlternativesOp : TransformDialectOp<"alternatives",
         ["getSuccessorEntryOperands", "getSuccessorRegions",
          "getRegionInvocationBounds"]>,
      DeclareOpInterfaceMethods<TransformOpInterface>,
-     FunctionalStyleTransformOpTrait, MemoryEffectsOpInterface,
+     DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
      IsolatedFromAbove, PossibleTopLevelTransformOpTrait,
      SingleBlockImplicitTerminator<"::mlir::transform::YieldOp">]> {
   let summary = "Attempts sequences of transforms until one succeeds";

diff  --git a/mlir/lib/Dialect/Transform/IR/TransformOps.cpp b/mlir/lib/Dialect/Transform/IR/TransformOps.cpp
index 70bff0c6d2c86..8bd5cab2f460b 100644
--- a/mlir/lib/Dialect/Transform/IR/TransformOps.cpp
+++ b/mlir/lib/Dialect/Transform/IR/TransformOps.cpp
@@ -240,6 +240,17 @@ transform::AlternativesOp::apply(transform::TransformResults &results,
   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();

diff  --git a/mlir/lib/Dialect/Transform/Transforms/CheckUses.cpp b/mlir/lib/Dialect/Transform/Transforms/CheckUses.cpp
index 94a3c8a6c6fd8..05fba01d689cb 100644
--- a/mlir/lib/Dialect/Transform/Transforms/CheckUses.cpp
+++ b/mlir/lib/Dialect/Transform/Transforms/CheckUses.cpp
@@ -140,7 +140,13 @@ class TransformOpMemFreeAnalysis {
       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()
@@ -149,7 +155,8 @@ class TransformOpMemFreeAnalysis {
     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

diff  --git a/mlir/test/Dialect/Transform/check-use-after-free.mlir b/mlir/test/Dialect/Transform/check-use-after-free.mlir
index a461d2548d5bb..80761639edd12 100644
--- a/mlir/test/Dialect/Transform/check-use-after-free.mlir
+++ b/mlir/test/Dialect/Transform/check-use-after-free.mlir
@@ -167,3 +167,13 @@ func.func @use_after_free_cycle() {
   return
 }
 
+// -----
+
+// This should not crash.
+
+transform.sequence failures(propagate) {
+^bb0(%arg0: !pdl.operation):
+  alternatives %arg0 : !pdl.operation {
+  ^bb0(%arg1: !pdl.operation):
+  }
+}