Reapply "Extend getBackwardSlice to track values captured…

[IanWood1]

This commit fixes the failure in the original PR when building with shared libs. The problem is that visitUsedValuesDefinedAbove is defined in MLIRTransformUtils, but that lib depends on this lib (MLIRAnalysis). To fix, I dropped the use of visitUsedValuesDefinedAbove and use Region::walk to traverse values defined above.

The 1st commit reapplies the original commit, and the 2nd fixes the build failure.

Reapplies #113478
Reverts #114432

This reverts commit a9a8351.

[llvmbot]

@llvm/pr-subscribers-mlir

Author: Ian Wood (IanWood1)

Changes

This commit fixes the failure in the original PR when building with shared libs. The problem is that visitUsedValuesDefinedAbove is defined in MLIRTransformUtils, but that lib depends on this lib (MLIRAnalysis). To fix, I dropped the use of visitUsedValuesDefinedAbove and use Region::walk to traverse values defined above.

The 1st commit reapplies the original commit, and the 2nd fixes the build failure.

Reapplies #113478
Reverts #114432

This reverts commit a9a8351.

---

Full diff: https://github.com/llvm/llvm-project/pull/114452.diff

4 Files Affected:

• (modified) mlir/include/mlir/Analysis/SliceAnalysis.h (+5)
• (modified) mlir/lib/Analysis/SliceAnalysis.cpp (+22-5)
• (modified) mlir/test/IR/slice.mlir (+27-1)
• (modified) mlir/test/lib/IR/TestSlicing.cpp (+2)

diff --git a/mlir/include/mlir/Analysis/SliceAnalysis.h b/mlir/include/mlir/Analysis/SliceAnalysis.h
index 99279fdfe427c8..a4f5d937cd51da 100644
--- a/mlir/include/mlir/Analysis/SliceAnalysis.h
+++ b/mlir/include/mlir/Analysis/SliceAnalysis.h
@@ -47,6 +47,11 @@ struct BackwardSliceOptions : public SliceOptions {
   /// backward slice computation traverses block arguments and asserts that the
   /// parent op has a single region with a single block.
   bool omitBlockArguments = false;
+
+  /// When omitUsesFromAbove is true, the backward slice computation omits
+  /// traversing values that are captured from above.
+  /// TODO: this should default to `false` after users have been updated.
+  bool omitUsesFromAbove = true;
 };
 
 using ForwardSliceOptions = SliceOptions;
diff --git a/mlir/lib/Analysis/SliceAnalysis.cpp b/mlir/lib/Analysis/SliceAnalysis.cpp
index 2b1cf411ceeeeb..8803ba994b2c14 100644
--- a/mlir/lib/Analysis/SliceAnalysis.cpp
+++ b/mlir/lib/Analysis/SliceAnalysis.cpp
@@ -16,6 +16,7 @@
 #include "mlir/IR/Operation.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"
 #include "mlir/Support/LLVM.h"
+#include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/SmallPtrSet.h"
 
@@ -91,14 +92,13 @@ static void getBackwardSliceImpl(Operation *op,
   if (options.filter && !options.filter(op))
     return;
 
-  for (const auto &en : llvm::enumerate(op->getOperands())) {
-    auto operand = en.value();
-    if (auto *definingOp = operand.getDefiningOp()) {
+  auto processValue = [&](Value value) {
+    if (auto *definingOp = value.getDefiningOp()) {
       if (backwardSlice->count(definingOp) == 0)
         getBackwardSliceImpl(definingOp, backwardSlice, options);
-    } else if (auto blockArg = dyn_cast<BlockArgument>(operand)) {
+    } else if (auto blockArg = dyn_cast<BlockArgument>(value)) {
       if (options.omitBlockArguments)
-        continue;
+        return;
 
       Block *block = blockArg.getOwner();
       Operation *parentOp = block->getParentOp();
@@ -113,7 +113,24 @@ static void getBackwardSliceImpl(Operation *op,
     } else {
       llvm_unreachable("No definingOp and not a block argument.");
     }
+  };
+
+  if (!options.omitUsesFromAbove) {
+    llvm::for_each(op->getRegions(), [&](Region &region) {
+      // Walk this region recursively to collect the regions that descend from
+      // this op's nested regions (inclusive).
+      SmallPtrSet<Region *, 4> descendents;
+      region.walk(
+          [&](Region *childRegion) { descendents.insert(childRegion); });
+      region.walk([&](Operation *op) {
+        for (OpOperand &operand : op->getOpOperands()) {
+          if (!descendents.contains(operand.get().getParentRegion()))
+            processValue(operand.get());
+        }
+      });
+    });
   }
+  llvm::for_each(op->getOperands(), processValue);
 
   backwardSlice->insert(op);
 }
diff --git a/mlir/test/IR/slice.mlir b/mlir/test/IR/slice.mlir
index 0a32a0f231baf2..87d446c8f415af 100644
--- a/mlir/test/IR/slice.mlir
+++ b/mlir/test/IR/slice.mlir
@@ -1,4 +1,4 @@
-// RUN: mlir-opt -slice-analysis-test %s | FileCheck %s
+// RUN: mlir-opt -slice-analysis-test -split-input-file %s | FileCheck %s
 
 func.func @slicing_linalg_op(%arg0 : index, %arg1 : index, %arg2 : index) {
   %a = memref.alloc(%arg0, %arg2) : memref<?x?xf32>
@@ -33,3 +33,29 @@ func.func @slicing_linalg_op(%arg0 : index, %arg1 : index, %arg2 : index) {
 //   CHECK-DAG:   %[[B:.+]] = memref.alloc(%[[ARG2]], %[[ARG1]]) : memref<?x?xf32>
 //   CHECK-DAG:   %[[C:.+]] = memref.alloc(%[[ARG0]], %[[ARG1]]) : memref<?x?xf32>
 //       CHECK:   return
+
+// -----
+
+#map = affine_map<(d0, d1) -> (d0, d1)>
+func.func @slice_use_from_above(%arg0: tensor<5x5xf32>, %arg1: tensor<5x5xf32>) {
+  %0 = linalg.generic {indexing_maps = [#map, #map], iterator_types = ["parallel", "parallel"]} ins(%arg0 : tensor<5x5xf32>) outs(%arg1 : tensor<5x5xf32>) {
+  ^bb0(%in: f32, %out: f32):
+    %2 = arith.addf %in, %in : f32
+    linalg.yield %2 : f32
+  } -> tensor<5x5xf32>
+  %collapsed = tensor.collapse_shape %0 [[0, 1]] : tensor<5x5xf32> into tensor<25xf32>
+  %1 = linalg.generic {indexing_maps = [#map, #map], iterator_types = ["parallel", "parallel"]} ins(%0 : tensor<5x5xf32>) outs(%arg1 : tensor<5x5xf32>) {
+  ^bb0(%in: f32, %out: f32):
+    %c2 = arith.constant 2 : index
+    %extracted = tensor.extract %collapsed[%c2] : tensor<25xf32>
+    %2 = arith.addf %extracted, %extracted : f32
+    linalg.yield %2 : f32
+  } -> tensor<5x5xf32>
+  return
+}
+
+// CHECK-LABEL: func @slice_use_from_above__backward_slice__0
+//  CHECK-SAME:   %[[ARG0:[a-zA-Z0-9_]+]]: tensor 
+//       CHECK:   %[[A:.+]] = linalg.generic {{.*}} ins(%[[ARG0]]
+//       CHECK:   %[[B:.+]] = tensor.collapse_shape %[[A]]
+//       CHECK:   return
diff --git a/mlir/test/lib/IR/TestSlicing.cpp b/mlir/test/lib/IR/TestSlicing.cpp
index c3d0d151c6d24d..e99d5976d6d9df 100644
--- a/mlir/test/lib/IR/TestSlicing.cpp
+++ b/mlir/test/lib/IR/TestSlicing.cpp
@@ -39,6 +39,8 @@ static LogicalResult createBackwardSliceFunction(Operation *op,
   SetVector<Operation *> slice;
   BackwardSliceOptions options;
   options.omitBlockArguments = omitBlockArguments;
+  // TODO: Make this default.
+  options.omitUsesFromAbove = false;
   getBackwardSlice(op, &slice, options);
   for (Operation *slicedOp : slice)
     builder.clone(*slicedOp, mapper);