Extend getBackwardSlice to track values captured from above

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;

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 "mlir/Transforms/RegionUtils.h"
 #include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/SmallPtrSet.h"
 
@@ -115,6 +116,19 @@ static void getBackwardSliceImpl(Operation *op,
     }
   }
 
+  // Visit values that are defined above.
+  if (!options.omitUsesFromAbove) {
+    visitUsedValuesDefinedAbove(op->getRegions(), [&](OpOperand *operand) {
+      if (Operation *definingOp = operand->get().getDefiningOp()) {
+        getBackwardSliceImpl(definingOp, backwardSlice, options);
+        return;
+      }
+      Operation *bbAargOwner =
+          cast<BlockArgument>(operand->get()).getOwner()->getParentOp();
+      getBackwardSliceImpl(bbAargOwner, backwardSlice, options);
+    });
+  }
+
   backwardSlice->insert(op);
 }