[MLIR] getBackwardSlice: don't bail on ops that are IsolatedFromAbove (#158135)

Ops with the `IsIsolatedFromAbove` trait should be captured by the
backward slice.

---------

Signed-off-by: Ian Wood <[email protected]>

Author: IanWood1

mlir/lib/Analysis/SliceAnalysis.cpp

@@ -109,7 +109,7 @@ static LogicalResult getBackwardSliceImpl(Operation *op,
                                           DenseSet<Operation *> &visited,
                                           SetVector<Operation *> *backwardSlice,
                                           const BackwardSliceOptions &options) {
-  if (!op || op->hasTrait<OpTrait::IsIsolatedFromAbove>())
+  if (!op)
     return success();
 
   // Evaluate whether we should keep this def.
@@ -136,7 +136,8 @@ static LogicalResult getBackwardSliceImpl(Operation *op,
       // blocks of parentOp, which are not technically backward unless they flow
       // into us. For now, just bail.
       if (parentOp && backwardSlice->count(parentOp) == 0) {
-        if (parentOp->getNumRegions() == 1 &&
+        if (!parentOp->hasTrait<OpTrait::IsIsolatedFromAbove>() &&
+            parentOp->getNumRegions() == 1 &&
             parentOp->getRegion(0).hasOneBlock()) {
           return getBackwardSliceImpl(parentOp, visited, backwardSlice,
                                       options);
@@ -150,7 +151,8 @@ static LogicalResult getBackwardSliceImpl(Operation *op,
 
   bool succeeded = true;
 
-  if (!options.omitUsesFromAbove) {
+  if (!options.omitUsesFromAbove &&
+      !op->hasTrait<OpTrait::IsIsolatedFromAbove>()) {
     llvm::for_each(op->getRegions(), [&](Region &region) {
       // Walk this region recursively to collect the regions that descend from
       // this op's nested regions (inclusive).

mlir/test/Transforms/move-operation-deps.mlir

@@ -460,3 +460,31 @@ module attributes {transform.with_named_sequence} {
     transform.yield
   }
 }
+
+// -----
+
+func.func @move_isolated_from_above() -> () {
+  %1 = "before"() : () -> (f32)
+  %2 = "moved0"() : () -> (f32)
+  %3 = test.isolated_one_region_op %2 {} : f32 -> f32
+  %4 = "moved1"(%3) : (f32) -> (f32)
+  return
+}
+// CHECK-LABEL: func @move_isolated_from_above()
+//       CHECK:   %[[MOVED0:.+]] = "moved0"
+//       CHECK:   %[[ISOLATED:.+]] = test.isolated_one_region_op %[[MOVED0]]
+//       CHECK:   %[[MOVED1:.+]] = "moved1"(%[[ISOLATED]])
+//       CHECK:   %[[BEFORE:.+]] = "before"
+
+module attributes {transform.with_named_sequence} {
+  transform.named_sequence @__transform_main(%arg0 : !transform.any_op {transform.readonly}) {
+    %op1 = transform.structured.match ops{["before"]} in %arg0
+        : (!transform.any_op) -> !transform.any_op
+    %op2 = transform.structured.match ops{["moved1"]} in %arg0
+        : (!transform.any_op) -> !transform.any_op
+    %v1 = transform.get_result %op2[0] : (!transform.any_op) -> !transform.any_value
+    transform.test.move_value_defns %v1 before %op1
+        : (!transform.any_value), !transform.any_op
+    transform.yield
+  }
+}

mlir/test/lib/Dialect/Test/TestOps.td

@@ -552,9 +552,10 @@ def OneRegionWithOperandsOp : TEST_Op<"one_region_with_operands_op", []> {
 
 def IsolatedOneRegionOp : TEST_Op<"isolated_one_region_op", [IsolatedFromAbove]> {
   let arguments = (ins Variadic<AnyType>:$operands);
+  let results = (outs Variadic<AnyType>:$results);
   let regions = (region AnyRegion:$my_region);
   let assemblyFormat = [{
-    attr-dict-with-keyword $operands $my_region `:` type($operands)
+    attr-dict-with-keyword $operands $my_region `:` type($operands) `->` type($results)
   }];
 }
 

mlir/test/lib/Transforms/TestMakeIsolatedFromAbove.cpp

@@ -27,8 +27,8 @@ makeIsolatedFromAboveImpl(RewriterBase &rewriter,
       makeRegionIsolatedFromAbove(rewriter, region, callBack);
   SmallVector<Value> operands = regionOp.getOperands();
   operands.append(capturedValues);
-  auto isolatedRegionOp =
-      test::IsolatedOneRegionOp::create(rewriter, regionOp.getLoc(), operands);
+  auto isolatedRegionOp = test::IsolatedOneRegionOp::create(
+      rewriter, regionOp.getLoc(), TypeRange(), operands);
   rewriter.inlineRegionBefore(region, isolatedRegionOp.getRegion(),
                               isolatedRegionOp.getRegion().begin());
   rewriter.eraseOp(regionOp);