turn into a general slice walk

Author: Dinistro

mlir/include/mlir/IR/SliceSupport.h

@@ -15,16 +15,13 @@ namespace mlir {
 
 /// A class to signal how to proceed with the walk of the backward slice:
 /// - Interrupt: Stops the walk.
-/// - Advance: Continues the walk to control flow predecessors values.
 /// - AdvanceTo: Continues the walk to user-specified values.
 /// - Skip: Continues the walk, but skips the predecessors of the current value.
 class WalkContinuation {
 public:
   enum class WalkAction {
     /// Stops the walk.
     Interrupt,
-    /// Continues the walk to control flow predecessors values.
-    Advance,
     /// Continues the walk to user-specified values.
     AdvanceTo,
     /// Continues the walk, but skips the predecessors of the current value.
@@ -34,10 +31,6 @@ class WalkContinuation {
   WalkContinuation(WalkAction action, mlir::ValueRange nextValues)
       : action(action), nextValues(nextValues) {}
 
-  /// Allows LogicalResult to interrupt the walk on failure.
-  explicit WalkContinuation(llvm::LogicalResult action)
-      : action(failed(action) ? WalkAction::Interrupt : WalkAction::Advance) {}
-
   /// Allows diagnostics to interrupt the walk.
   explicit WalkContinuation(mlir::Diagnostic &&)
       : action(WalkAction::Interrupt) {}
@@ -58,12 +51,6 @@ class WalkContinuation {
     return WalkContinuation(WalkAction::AdvanceTo, nextValues);
   }
 
-  /// Creates a continuation that adds the control flow predecessor values to
-  /// the work list and advances the walk.
-  static WalkContinuation advance() {
-    return WalkContinuation(WalkAction::Advance, {});
-  }
-
   /// Creates a continuation that advances the walk without adding any
   /// predecessor values to the work list.
   static WalkContinuation skip() {
@@ -89,34 +76,23 @@ class WalkContinuation {
 };
 
 /// A callback that is invoked for each value encountered during the walk of the
-/// backward slice. The callback takes the current value, and returns the walk
-/// continuation, which determines if the walk should proceed and if yes, with
-/// which values.
-using WalkCallback = mlir::function_ref<WalkContinuation(mlir::Value)>;
-
-/// Walks the backward slice starting from the `rootValues` using a depth-first
-/// traversal following the use-def chains. The walk calls the provided
-/// `walkCallback` for each value encountered in the backward slice and uses the
-/// returned walk continuation to determine how to proceed. Additionally, the
-/// walk also transparently traverses through select operations and control flow
-/// operations that implement RegionBranchOpInterface or BranchOpInterface.
-WalkContinuation walkBackwardSlice(mlir::ValueRange rootValues,
-                                   WalkCallback walkCallback);
-
-/// A callback that is invoked for each value encountered during the walk of the
-/// backward slice. The callback takes the current value, and returns the walk
+/// slice. The callback takes the current value, and returns the walk
 /// continuation, which determines if the walk should proceed and if yes, with
 /// which values.
 using WalkCallback = mlir::function_ref<WalkContinuation(mlir::Value)>;
 
-/// Walks the backward slice starting from the `rootValues` using a depth-first
-/// traversal following the use-def chains. The walk calls the provided
-/// `walkCallback` for each value encountered in the backward slice and uses the
-/// returned walk continuation to determine how to proceed. Additionally, the
-/// walk also transparently traverses through select operations and control flow
-/// operations that implement RegionBranchOpInterface or BranchOpInterface.
-WalkContinuation walkBackwardSlice(mlir::ValueRange rootValues,
-                                   WalkCallback walkCallback);
+/// Walks the slice starting from the `rootValues` using a depth-first
+/// traversal. The walk calls the provided `walkCallback` for each value
+/// encountered in the slice and uses the returned walk continuation to
+/// determine how to proceed.
+WalkContinuation walkSlice(mlir::ValueRange rootValues,
+                           WalkCallback walkCallback);
+
+/// Computes a vector of all control predecessors of `value`. Relies on
+/// RegionBranchOpInterface and BranchOpInterface to determine predecessors.
+/// Returns nullopt if value has no predecessors or when the relevant operations
+/// are missing the interface implementations.
+std::optional<SmallVector<Value>> getControlFlowPredecessors(Value value);
 
 } // namespace mlir
 

mlir/lib/Dialect/LLVMIR/Transforms/InlinerInterfaceImpl.cpp

@@ -224,10 +224,11 @@ static ArrayAttr concatArrayAttr(ArrayAttr lhs, ArrayAttr rhs) {
 
 /// Attempts to return the set of all underlying pointer values that
 /// `pointerValue` is based on. This function traverses through select
-/// operations and block arguments unlike getUnderlyingObject.
-static SmallVector<Value> getUnderlyingObjectSet(Value pointerValue) {
+/// operations and block arguments.
+static FailureOr<SmallVector<Value>>
+getUnderlyingObjectSet(Value pointerValue) {
   SmallVector<Value> result;
-  walkBackwardSlice(pointerValue, [&](Value val) {
+  WalkContinuation walkResult = walkSlice(pointerValue, [&](Value val) {
     if (auto gepOp = val.getDefiningOp<LLVM::GEPOp>())
       return WalkContinuation::advanceTo(gepOp.getBase());
 
@@ -239,12 +240,28 @@ static SmallVector<Value> getUnderlyingObjectSet(Value pointerValue) {
       return WalkContinuation::advanceTo(
           {selectOp.getTrueValue(), selectOp.getFalseValue()});
 
-    if (isa<OpResult>(val))
+    // Attempt to advance to control flow predecessors.
+    std::optional<SmallVector<Value>> controlFlowPredecessors =
+        getControlFlowPredecessors(val);
+    if (controlFlowPredecessors)
+      return WalkContinuation::advanceTo(*controlFlowPredecessors);
+
+    // For all non-control flow results, consider `val` an underlying object.
+    if (isa<OpResult>(val)) {
       result.push_back(val);
+      return WalkContinuation::skip();
+    }
 
-    return WalkContinuation::advance();
+    // If this place is reached, `val` is a block argument that is not
+    // understood. Therfore, we conservatively interrupt.
+    // Note: Dealing with the function arguments is not necessary, as the slice
+    // would have to go through an SSACopyOp first.
+    return WalkContinuation::interrupt();
   });
 
+  if (walkResult.wasInterrupted())
+    return failure();
+
   return result;
 }
 
@@ -306,9 +323,14 @@ static void createNewAliasScopesFromNoAliasParameter(
 
       // Find the set of underlying pointers that this pointer is based on.
       SmallPtrSet<Value, 4> basedOnPointers;
-      for (Value pointer : pointerArgs)
-        llvm::copy(getUnderlyingObjectSet(pointer),
+      for (Value pointer : pointerArgs) {
+        FailureOr<SmallVector<Value>> underlyingObjectSet =
+            getUnderlyingObjectSet(pointer);
+        if (failed(underlyingObjectSet))
+          return;
+        llvm::copy(*underlyingObjectSet,
                    std::inserter(basedOnPointers, basedOnPointers.begin()));
+      }
 
       bool aliasesOtherKnownObject = false;
       // Go through the based on pointers and check that they are either:

mlir/lib/IR/SliceSupport.cpp

@@ -3,6 +3,35 @@
 
 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>
@@ -49,15 +78,20 @@ getRegionPredecessorOperands(RegionBranchOpInterface regionOp,
   return predecessorOperands;
 }
 
-/// Returns the predecessor branch operands that match `blockArg`.
-static SmallVector<Value> getBlockPredecessorOperands(BlockArgument blockArg) {
+/// Returns the predecessor branch operands that match `blockArg`. Returns a
+/// nullopt when some of the predecessor terminators do not implement the
+/// BranchOpInterface.
+static std::optional<SmallVector<Value>>
+getBlockPredecessorOperands(BlockArgument blockArg) {
   Block *block = blockArg.getOwner();
 
   // Search the predecessor operands for all predecessor terminators.
   SmallVector<Value> predecessorOperands;
   for (auto it = block->pred_begin(); it != block->pred_end(); ++it) {
     Block *predecessor = *it;
-    auto branchOp = cast<BranchOpInterface>(predecessor->getTerminator());
+    auto branchOp = dyn_cast<BranchOpInterface>(predecessor->getTerminator());
+    if (!branchOp)
+      return std::nullopt;
     SuccessorOperands successorOperands =
         branchOp.getSuccessorOperands(it.getSuccessorIndex());
     // Store the predecessor operand if the block argument matches an operand
@@ -69,62 +103,38 @@ static SmallVector<Value> getBlockPredecessorOperands(BlockArgument blockArg) {
   return predecessorOperands;
 }
 
-mlir::WalkContinuation mlir::walkBackwardSlice(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;
-
-    // Add the next values to the work list if the walk should continue.
-    if (continuation.wasAdvancedTo()) {
-      workList.append(continuation.getNextValues().begin(),
-                      continuation.getNextValues().end());
-      continue;
-    }
-
+std::optional<SmallVector<Value>>
+mlir::getControlFlowPredecessors(Value value) {
+  SmallVector<Value> result;
+  if (OpResult opResult = dyn_cast<OpResult>(value)) {
+    auto regionOp = dyn_cast<RegionBranchOpInterface>(opResult.getOwner());
+    // If the interface is not implemented, there are no control flow
+    // predecessors to work with.
+    if (!regionOp)
+      return std::nullopt;
     // Add the control flow predecessor operands to the work list.
-    if (OpResult opResult = dyn_cast<OpResult>(current)) {
-      auto regionOp = dyn_cast<RegionBranchOpInterface>(opResult.getOwner());
-      if (!regionOp)
-        continue;
-      RegionSuccessor region(regionOp->getResults());
-      SmallVector<Value> predecessorOperands = getRegionPredecessorOperands(
-          regionOp, region, opResult.getResultNumber());
-      workList.append(predecessorOperands.begin(), predecessorOperands.end());
-      continue;
-    }
+    RegionSuccessor region(regionOp->getResults());
+    SmallVector<Value> predecessorOperands = getRegionPredecessorOperands(
+        regionOp, region, opResult.getResultNumber());
+    return predecessorOperands;
+  }
 
-    auto blockArg = cast<BlockArgument>(current);
-    Block *block = blockArg.getOwner();
-    // Search the region predecessor operands for structured control flow.
-    auto regionBranchOp =
-        dyn_cast<RegionBranchOpInterface>(block->getParentOp());
-    if (block->isEntryBlock() && regionBranchOp) {
+  auto blockArg = cast<BlockArgument>(value);
+  Block *block = blockArg.getOwner();
+  // Search the region predecessor operands for structured control flow.
+  if (block->isEntryBlock()) {
+    if (auto regionBranchOp =
+            dyn_cast<RegionBranchOpInterface>(block->getParentOp())) {
       RegionSuccessor region(blockArg.getParentRegion());
       SmallVector<Value> predecessorOperands = getRegionPredecessorOperands(
           regionBranchOp, region, blockArg.getArgNumber());
-      workList.append(predecessorOperands.begin(), predecessorOperands.end());
-      continue;
+      return predecessorOperands;
     }
-    // Search the block predecessor operands for unstructured control flow.
-    SmallVector<Value> predecessorOperands =
-        getBlockPredecessorOperands(blockArg);
-    workList.append(predecessorOperands.begin(), predecessorOperands.end());
+    // Unclear how to deal with this operation, conservatively return a
+    // failure.
+    return std::nullopt;
   }
 
-  return WalkContinuation::advance();
+  // Search the block predecessor operands for unstructured control flow.
+  return getBlockPredecessorOperands(blockArg);
 }