[Mlir-commits] [mlir] [MLIR, LLVM]: Add an IR utility to perform slice walking (PR #103053)

[Tobias Gysi]

================
@@ -0,0 +1,140 @@
+#include "mlir/IR/SliceSupport.h"
+#include "mlir/Interfaces/ControlFlowInterfaces.h"
+
+using namespace mlir;
+
+WalkContinuation mlir::walkSlice(ValueRange rootValues,
+                                 WalkCallback walkCallback) {
+  // Search the backward slice starting from the root values.
+  SmallVector<Value> workList = rootValues;
+  llvm::SmallDenseSet<Value, 16> seenValues;
+  while (!workList.empty()) {
+    // Search the backward slice of the current value.
+    Value current = workList.pop_back_val();
+
+    // Skip the current value if it has already been seen.
+    if (!seenValues.insert(current).second)
+      continue;
+
+    // Call the walk callback with the current value.
+    WalkContinuation continuation = walkCallback(current);
+    if (continuation.wasInterrupted())
+      return continuation;
+    if (continuation.wasSkipped())
+      continue;
+
+    assert(continuation.wasAdvancedTo());
+    // Add the next values to the work list if the walk should continue.
+    workList.append(continuation.getNextValues().begin(),
+                    continuation.getNextValues().end());
+  }
+
+  return WalkContinuation::skip();
+}
+
+/// Returns the operands from all predecessor regions that match `operandNumber`
+/// for the `successor` region within `regionOp`.
+static SmallVector<Value>
+getRegionPredecessorOperands(RegionBranchOpInterface regionOp,
+                             RegionSuccessor successor,
+                             unsigned operandNumber) {
+  SmallVector<Value> predecessorOperands;
+
+  // Returns true if `successors` contains `successor`.
+  auto isContained = [](ArrayRef<RegionSuccessor> successors,
+                        RegionSuccessor successor) {
+    auto *it = llvm::find_if(successors, [&successor](RegionSuccessor curr) {
+      return curr.getSuccessor() == successor.getSuccessor();
+    });
+    return it != successors.end();
+  };
+
+  // Search the operand ranges on the region operation itself.
+  SmallVector<Attribute> operandAttributes(regionOp->getNumOperands());
+  SmallVector<RegionSuccessor> successors;
+  regionOp.getEntrySuccessorRegions(operandAttributes, successors);
+  if (isContained(successors, successor)) {
+    OperandRange operands = regionOp.getEntrySuccessorOperands(successor);
+    predecessorOperands.push_back(operands[operandNumber]);
+  }
+
+  // Search the operand ranges on region terminators.
+  for (Region &region : regionOp->getRegions()) {
+    for (Block &block : region) {
+      auto terminatorOp =
+          dyn_cast<RegionBranchTerminatorOpInterface>(block.getTerminator());
+      if (!terminatorOp)
+        continue;
+      SmallVector<Attribute> operandAttributes(terminatorOp->getNumOperands());
+      SmallVector<RegionSuccessor> successors;
+      terminatorOp.getSuccessorRegions(operandAttributes, successors);
+      if (isContained(successors, successor)) {
+        OperandRange operands = terminatorOp.getSuccessorOperands(successor);
+        predecessorOperands.push_back(operands[operandNumber]);
+      }
+    }
+  }
+
+  return predecessorOperands;
+}
+
+/// Returns the predecessor branch operands that match `blockArg`. Returns a
+/// nullopt when some of the predecessor terminators do not implement the
----------------
gysit wrote:

```suggestion
/// Returns the predecessor branch operands that match `blockArg`, or 
/// nullopt if some of the predecessor terminators do not implement the
```
nit:

https://github.com/llvm/llvm-project/pull/103053