[mlir][bufferization] Improve findValueInReverseUseDefChain signature

[springerm]

Summary

Instead of passing traversal options as a long list of arguments, store them in a TraversalConfig object and pass that object.

mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h

struct BufferizationOptions {
   /// Initializer functions for analysis state. These can be used to
   /// initialize dialect-specific analysis state.
   SmallVector<AnalysisStateInitFn> stateInitializers;
 };
 
 /// Return `true` if the given value is a BlockArgument of a func::FuncOp.
 bool isFunctionArgument(Value value);
 
+/// Traversal parameters for `findValueInReverseUseDefChain`.
+struct TraversalConfig {
+  /// Specifies if leaves (that do not have further OpOperands to follow)
+  /// should be returned even if they do not match the specified filter.
+  bool alwaysIncludeLeaves = true;
+
+  /// Specifies whether out-of-place/undecided OpOperands should be followed.
+  bool followInPlaceOnly = false;
+
+  /// Specifies whether non-equivalent OpOperands should be followed.
+  bool followEquivalentOnly = false;
+
+  /// Specifies whether unknown/non-bufferizable/ops not included in the
+  /// OpFilter of BufferizationOptions should be followed.
+  bool followUnknownOps = false;
+};
+
 /// AnalysisState provides a variety of helper functions for dealing with
 /// tensor values.
 class AnalysisState {
 public:
   /// Determine which OpOperand* will alias with `result` if the op is
   /// bufferized in place. Return all tensor OpOperand* if the op is not
   /// bufferizable.
   AliasingOpOperandList getAliasingOpOperands(OpResult result) const;
public:
   /// read by themselves (e.g., ExtractSliceOp).
   bool isValueRead(Value value) const;
 
   /// Starting from `value`, follow the use-def chain in reverse, always
   /// selecting the aliasing OpOperands. Find and return Values for which
   /// `condition` evaluates to true. OpOperands of such matching Values are not
   /// traversed any further.
   ///
-  /// When reaching the end of a chain (BlockArgument or Value without aliasing
-  /// OpOperands), also return the last Value of that chain if
-  /// `alwaysIncludeLeaves` is set.
+  /// When reaching the end of a chain, also return the last Value of that
+  /// chain if `config.alwaysIncludeLeaves` is set.
   ///
   /// Example:
   ///
   ///                               8
   ///                               |
   ///   6*         7*         +-----+----+
   ///   |          |          |          |
   ///   2*         3          4*         5
   ///   |          |          |          |
   ///   +----------+----------+----------+
   ///              |
   ///              1
   ///
   /// In the above example, Values with a star satisfy the condition. When
   /// starting the traversal from Value 1, the resulting SetVector is:
   /// { 2, 7, 8, 5 }
   ///
-  /// If `followEquivalentOnly` is set, only equivalent OpOperands are selected.
+  /// Additional stopping conditions for the traversal can be specified in
+  /// `config`.
   SetVector<Value> findValueInReverseUseDefChain(
       Value value, llvm::function_ref<bool(Value)> condition,
-      bool followEquivalentOnly = false, bool alwaysIncludeLeaves = true) const;
+      TraversalConfig config = TraversalConfig()) const;
 
   /// Find the values that may define the contents of the given value at
   /// runtime. A block argument is always a definition. An OpResult is a
   /// definition if it bufferizes to memory write. If it does not bufferize to
   /// a memory write but has aliasing operands, we continue the lookup on these
   /// values.
   ///
   /// Example: %r = tensor.insert %f into %t[%c0] : tensor<?xf32>

mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp

if (!bufferizableOp)
     return nullptr;
   if (!isOpAllowed(op))
     return nullptr;
   return bufferizableOp;
 }
 
 BufferizableOpInterface
 BufferizationOptions::dynCastBufferizableOp(Value value) const {
-  if (auto bufferizableOp = value.getDefiningOp<BufferizableOpInterface>())
-    if (isOpAllowed(bufferizableOp.getOperation()))
-      return bufferizableOp;
-  return nullptr;
+  return dynCastBufferizableOp(getOwnerOfValue(value));
 }
 
 void BufferizationOptions::setFunctionBoundaryTypeConversion(
     LayoutMapOption layoutMapOption) {
   functionArgTypeConverterFn = [=](TensorType tensorType, Attribute memorySpace,
                                    func::FuncOp funcOp,
                                    const BufferizationOptions &options) {
     if (layoutMapOption == LayoutMapOption::IdentityLayoutMap)

 }
 
 // Starting from `value`, follow the use-def chain in reverse, always selecting
 // the aliasing OpOperands. Find and return Values for which `condition`
 // evaluates to true. OpOperands of such matching Values are not traversed any
 // further.
 llvm::SetVector<Value> AnalysisState::findValueInReverseUseDefChain(
     Value value, llvm::function_ref<bool(Value)> condition,
-    bool followEquivalentOnly, bool alwaysIncludeLeaves) const {
+    TraversalConfig config) const {
   llvm::SetVector<Value> result, workingSet;
   workingSet.insert(value);
 
   while (!workingSet.empty()) {
     Value value = workingSet.pop_back_val();
     if (condition(value)) {
       result.insert(value);
       continue;
     }
 
     if (llvm::isa<BlockArgument>(value)) {
-      if (alwaysIncludeLeaves)
+      if (config.alwaysIncludeLeaves)
         result.insert(value);
       continue;
     }
 
     OpResult opResult = llvm::cast<OpResult>(value);
     BufferizableOpInterface bufferizableOp =
         options.dynCastBufferizableOp(opResult.getDefiningOp());
-    AliasingOpOperandList aliases = getAliasingOpOperands(opResult);
+    if (!config.followUnknownOps && !bufferizableOp) {
+      // Stop iterating if `followUnknownOps` is unset and the op is either
+      // not bufferizable or excluded in the OpFilter.
+      if (config.alwaysIncludeLeaves)
+        result.insert(value);
+      continue;
+    }
 
-    // Stop iterating in either one of these cases:
-    // * The current op is not bufferizable or excluded in the filter.
-    // * There are no OpOperands to follow.
-    if (!bufferizableOp || aliases.getNumAliases() == 0) {
-      if (alwaysIncludeLeaves)
+    AliasingOpOperandList aliases = getAliasingOpOperands(opResult);
+    if (aliases.getNumAliases() == 0) {
+      // The traversal ends naturally if there are no more OpOperands that
+      // could be followed.
+      if (config.alwaysIncludeLeaves)
         result.insert(value);
       continue;
     }
 
     for (AliasingOpOperand a : aliases) {
-      if (followEquivalentOnly && a.relation != BufferRelation::Equivalent) {
+      if (config.followEquivalentOnly &&
+          a.relation != BufferRelation::Equivalent) {
         // Stop iterating if `followEquivalentOnly` is set but the alias is not
         // equivalent.
-        if (alwaysIncludeLeaves)
+        if (config.alwaysIncludeLeaves)
           result.insert(value);
       } else {
         workingSet.insert(a.opOperand->get());
       }
+
+      if (config.followInPlaceOnly && !isInPlace(*a.opOperand)) {
+        // Stop iterating if `followInPlaceOnly` is set but the alias is
+        // out-of-place.
+        if (config.alwaysIncludeLeaves)
+          result.insert(value);
+        continue;
+      }
+
+      workingSet.insert(a.opOperand->get());
     }
   }
 
   return result;
 }
 
 // Find the values that define the contents of the given value.
 llvm::SetVector<Value> AnalysisState::findDefinitions(Value value) const {
+  TraversalConfig config;
+  config.alwaysIncludeLeaves = false;
   return findValueInReverseUseDefChain(
-      value, [&](Value v) { return this->bufferizesToMemoryWrite(v); },
-      /*followEquivalentOnly=*/false, /*alwaysIncludeLeaves=*/false);
+      value, [&](Value v) { return this->bufferizesToMemoryWrite(v); }, config);
 }
 
 AnalysisState::AnalysisState(const BufferizationOptions &options)
     : AnalysisState(options, TypeID::get<AnalysisState>()) {}
 
 AnalysisState::AnalysisState(const BufferizationOptions &options, TypeID type)
     : options(options), type(type) {
   for (const BufferizationOptions::AnalysisStateInitFn &fn :
bool bufferization::detail::defaultResultBufferizesToMemoryWrite(
   // * conflictingWrite = %1
   //
   auto isMemoryWriteInsideOp = [&](Value v) {
     Operation *op = getOwnerOfValue(v);
     if (!opResult.getDefiningOp()->isAncestor(op))
       return false;
     return state.bufferizesToMemoryWrite(v);
   };
+  TraversalConfig config;
+  config.alwaysIncludeLeaves = false;
   for (AliasingOpOperand alias : opOperands) {
     if (!state
              .findValueInReverseUseDefChain(alias.opOperand->get(),
-                                            isMemoryWriteInsideOp,
-                                            /*followEquivalentOnly=*/false,
-                                            /*alwaysIncludeLeaves=*/false)
+                                            isMemoryWriteInsideOp, config)
              .empty())
       return true;
   }
   return false;
 }
 
 // Compute the AliasingOpOperandList for a given OpResult based on
 // getAliasingOpResults.

mlir/lib/Dialect/Bufferization/Transforms/EmptyTensorElimination.cpp

for (OpOperand &operand : op->getOpOperands()) {
       SmallVector<Value> neededValues;
       // Is this an anchor?
       if (!anchorMatchFunc(operand, neededValues))
         continue;
 
       // Find tensor.empty ops on the reverse SSA use-def chain. Only follow
       // equivalent tensors. I.e., stop when there are ops such as extract_slice
       // on the path.
+      TraversalConfig config;
+      config.followEquivalentOnly = true;
+      config.alwaysIncludeLeaves = false;
       SetVector<Value> emptyTensors = state.findValueInReverseUseDefChain(
           operand.get(), /*condition=*/
           [&](Value val) { return val.getDefiningOp<tensor::EmptyOp>(); },
-          /*followEquivalentOnly=*/true, /*alwaysIncludeLeaves=*/false);
+          config);
 
       for (Value v : emptyTensors) {
         Operation *emptyTensorOp = v.getDefiningOp();
 
         // Replace only if the types match. We do not support slices or casts.
         // TODO: This could be extended to support IR such as:
         // %0 = tensor.empty() : tensor<128xf32>
         // %1 = "some_op"(%0) : (tensor<128xf32>) -> (tensor<128xf32>)

mlir/lib/Dialect/Bufferization/Transforms/OneShotAnalysis.cpp

 
 //===----------------------------------------------------------------------===//
 // Bufferization analyses.
 //===----------------------------------------------------------------------===//
 
 // Find the values that define the contents of the given value.
 const llvm::SetVector<Value> &
 OneShotAnalysisState::findDefinitionsCached(Value value) {
-  if (!cachedDefinitions.count(value)) {
-    cachedDefinitions[value] = findValueInReverseUseDefChain(
-        value, [&](Value v) { return this->bufferizesToMemoryWrite(v); },
-        /*followEquivalentOnly=*/false, /*alwaysIncludeLeaves=*/false);
-  }
+  if (!cachedDefinitions.count(value))
+    cachedDefinitions[value] = findDefinitions(value);
   return cachedDefinitions[value];
 }
 
 void OneShotAnalysisState::resetCache() { cachedDefinitions.clear(); }
 
 /// Determine if `operand` can be bufferized in-place.
 static LogicalResult
 bufferizableInPlaceAnalysisImpl(OpOperand &operand, OneShotAnalysisState &state,