[Mlir-commits] [mlir] Extend `getBackwardSlice` to track values captured from above (PR #113478)

[Ian Wood]

https://github.com/IanWood1 created https://github.com/llvm/llvm-project/pull/113478

This change modifies `getBackwardSlice`  to track values captures by the regions of each operation that it traverses. Ignoring values captured from a parent region may lead to an incomplete program slice. However, there seems to be logic that depends on not traversing captured values, so this change preserves the default behavior by hiding this logic behind the `omitUsesFromAbove` flag.

>From 29209330d68e877cb92eca9f48b782244a97ecf9 Mon Sep 17 00:00:00 2001
From: Ian Wood <ianwood2024 at u.northwestern.edu>
Date: Wed, 23 Oct 2024 16:46:02 +0000
Subject: [PATCH] Add `omitUsesFromAbove` to getBackwardsSlice

`getBackwardsSlice` should track values captured by each op's region
that it traverses, and follow those defs.

However, there might be logic that depends on not traversing captured
values so this change preserves the default behavior by hiding this
logic behind the `omitUsesFromAbove` flag.
---
 mlir/include/mlir/Analysis/SliceAnalysis.h |  5 +++++
 mlir/lib/Analysis/SliceAnalysis.cpp        | 14 ++++++++++++++
 2 files changed, 19 insertions(+)

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..d07ae7b3ffa2c7 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 "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);
 }