[MLIR][LLVM]: Add an IR utility to perform slice walking

mlir/include/mlir/Analysis/SliceWalk.h

@@ -0,0 +1,98 @@
+//===- SliceWalk.h - Helpers for performing IR slice walks ---*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_ANALYSIS_SLICEWALK_H
+#define MLIR_ANALYSIS_SLICEWALK_H
+
+#include "mlir/IR/ValueRange.h"
+
+namespace mlir {
+
+/// A class to signal how to proceed with the walk of the backward slice:
+/// - Interrupt: Stops the walk.
+/// - 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 user-specified values.
+    AdvanceTo,
+    /// Continues the walk, but skips the predecessors of the current value.
+    Skip
+  };
+
+  WalkContinuation(WalkAction action, mlir::ValueRange nextValues)
+      : action(action), nextValues(nextValues) {}
+
+  /// Allows diagnostics to interrupt the walk.
+  explicit WalkContinuation(mlir::Diagnostic &&)
+      : action(WalkAction::Interrupt) {}
+
+  /// Allows diagnostics to interrupt the walk.
+  explicit WalkContinuation(mlir::InFlightDiagnostic &&)
+      : action(WalkAction::Interrupt) {}
+
+  /// Creates a continuation that interrupts the walk.
+  static WalkContinuation interrupt() {
+    return WalkContinuation(WalkAction::Interrupt, {});
+  }
+
+  /// Creates a continuation that adds the user-specified `nextValues` to the
+  /// work list and advances the walk.
+  static WalkContinuation advanceTo(mlir::ValueRange nextValues) {
+    return WalkContinuation(WalkAction::AdvanceTo, nextValues);
+  }
+
+  /// Creates a continuation that advances the walk without adding any
+  /// predecessor values to the work list.
+  static WalkContinuation skip() {
+    return WalkContinuation(WalkAction::Skip, {});
+  }
+
+  /// Returns true if the walk was interrupted.
+  bool wasInterrupted() const { return action == WalkAction::Interrupt; }
+
+  /// Returns true if the walk was skipped.
+  bool wasSkipped() const { return action == WalkAction::Skip; }
+
+  /// Returns true if the walk was advanced to user-specified values.
+  bool wasAdvancedTo() const { return action == WalkAction::AdvanceTo; }
+
+  /// Returns the next values to continue the walk with.
+  mlir::ArrayRef<mlir::Value> getNextValues() const { return nextValues; }
+
+private:
+  WalkAction action;
+  /// The next values to continue the walk with.
+  mlir::SmallVector<mlir::Value> nextValues;
+};
+
+/// A callback that is invoked for each value encountered during the walk of the
+/// 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 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
+
+#endif // MLIR_ANALYSIS_SLICEWALK_H

mlir/lib/Analysis/CMakeLists.txt

@@ -29,6 +29,7 @@ add_mlir_library(MLIRAnalysis
   Liveness.cpp
   CFGLoopInfo.cpp
   SliceAnalysis.cpp
+  SliceWalk.cpp
   TopologicalSortUtils.cpp
 
   AliasAnalysis/LocalAliasAnalysis.cpp

mlir/lib/Analysis/SliceWalk.cpp

@@ -0,0 +1,139 @@
+#include "mlir/Analysis/SliceWalk.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`, or nullopt if
+/// 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 = 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
+    // and is not produced by the terminator.
+    if (Value operand = successorOperands[blockArg.getArgNumber()])
+      predecessorOperands.push_back(operand);
+  }
+
+  return predecessorOperands;
+}
+
+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.
+    RegionSuccessor region(regionOp->getResults());
+    SmallVector<Value> predecessorOperands = getRegionPredecessorOperands(
+        regionOp, region, opResult.getResultNumber());
+    return predecessorOperands;
+  }
+
+  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());
+      return predecessorOperands;
+    }
+    // If the interface is not implemented, there are no control flow
+    // predecessors to work with.
+    return std::nullopt;
+  }
+
+  // Search the block predecessor operands for unstructured control flow.
+  return getBlockPredecessorOperands(blockArg);
+}

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

@@ -12,6 +12,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/LLVMIR/Transforms/InlinerInterfaceImpl.h"
+#include "mlir/Analysis/SliceWalk.h"
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "mlir/IR/Matchers.h"
 #include "mlir/Interfaces/DataLayoutInterfaces.h"
@@ -221,86 +222,45 @@ static ArrayAttr concatArrayAttr(ArrayAttr lhs, ArrayAttr rhs) {
   return ArrayAttr::get(lhs.getContext(), result);
 }
 
-/// Attempts to return the underlying pointer value that `pointerValue` is based
-/// on. This traverses down the chain of operations to the last operation
-/// producing the base pointer and returns it. If it encounters an operation it
-/// cannot further traverse through, returns the operation's result.
-static Value getUnderlyingObject(Value pointerValue) {
-  while (true) {
-    if (auto gepOp = pointerValue.getDefiningOp<LLVM::GEPOp>()) {
-      pointerValue = gepOp.getBase();
-      continue;
-    }
-
-    if (auto addrCast = pointerValue.getDefiningOp<LLVM::AddrSpaceCastOp>()) {
-      pointerValue = addrCast.getOperand();
-      continue;
-    }
-
-    break;
-  }
-
-  return pointerValue;
-}
-
 /// 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;
-
-  SmallVector<Value> workList{pointerValue};
-  // Avoid dataflow loops.
-  SmallPtrSet<Value, 4> seen;
-  do {
-    Value current = workList.pop_back_val();
-    current = getUnderlyingObject(current);
-
-    if (!seen.insert(current).second)
-      continue;
-
-    if (auto selectOp = current.getDefiningOp<LLVM::SelectOp>()) {
-      workList.push_back(selectOp.getTrueValue());
-      workList.push_back(selectOp.getFalseValue());
-      continue;
+  WalkContinuation walkResult = walkSlice(pointerValue, [&](Value val) {
+    if (auto gepOp = val.getDefiningOp<LLVM::GEPOp>())
+      return WalkContinuation::advanceTo(gepOp.getBase());
+
+    if (auto addrCast = val.getDefiningOp<LLVM::AddrSpaceCastOp>())
+      return WalkContinuation::advanceTo(addrCast.getOperand());
+
+    // TODO: Add a SelectLikeOpInterface and use it in the slicing utility.
+    if (auto selectOp = val.getDefiningOp<LLVM::SelectOp>())
+      return WalkContinuation::advanceTo(
+          {selectOp.getTrueValue(), selectOp.getFalseValue()});
+
+    // 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();
     }
 
-    if (auto blockArg = dyn_cast<BlockArgument>(current)) {
-      Block *parentBlock = blockArg.getParentBlock();
-
-      // Attempt to find all block argument operands for every predecessor.
-      // If any operand to the block argument wasn't found in a predecessor,
-      // conservatively add the block argument to the result set.
-      SmallVector<Value> operands;
-      bool anyUnknown = false;
-      for (auto iter = parentBlock->pred_begin();
-           iter != parentBlock->pred_end(); iter++) {
-        auto branch = dyn_cast<BranchOpInterface>((*iter)->getTerminator());
-        if (!branch) {
-          result.push_back(blockArg);
-          anyUnknown = true;
-          break;
-        }
-
-        Value operand = branch.getSuccessorOperands(
-            iter.getSuccessorIndex())[blockArg.getArgNumber()];
-        if (!operand) {
-          result.push_back(blockArg);
-          anyUnknown = true;
-          break;
-        }
-
-        operands.push_back(operand);
-      }
-
-      if (!anyUnknown)
-        llvm::append_range(workList, operands);
-
-      continue;
-    }
+    // If this place is reached, `val` is a block argument that is not
+    // understood. Therefore, we conservatively interrupt.
+    // Note: Dealing with function arguments is not necessary, as the slice
+    // would have to go through an SSACopyOp first.
+    return WalkContinuation::interrupt();
+  });
 
-    result.push_back(current);
-  } while (!workList.empty());
+  if (walkResult.wasInterrupted())
+    return failure();
 
   return result;
 }
@@ -363,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: