[mlir][bufferization] Improve aliasing OpOperand/OpResult property

[springerm]

Summary

getAliasingOpOperands/getAliasingOpResults now encodes OpOperand/OpResult, buffer relation and a degree of certainty. E.g.:

// aliasingOpOperands(%r) = {(%t, EQUIV, DEFINITE)}
// aliasingOpResults(%t) = {(%r, EQUIV, DEFINITE)}
%r = tensor.insert %f into %t[%idx] : tensor<?xf32>

// aliasingOpOperands(%r) = {(%t0, EQUIV, MAYBE), (%t1, EQUIV, MAYBE)}
// aliasingOpResults(%t0) = {(%r, EQUIV, MAYBE)}
// aliasingOpResults(%t1) = {(%r, EQUIV, MAYBE)}
%r = arith.select %c, %t0, %t1 : tensor<?xf32>

BufferizableOpInterface::bufferRelation is removed, as it is now part of getAliasingOpOperands/getAliasingOpResults.

This change allows for better analysis, in particular wrt. equivalence. This allows additional optimizations and better error checking (which is sometimes overly conservative). Examples:

• EmptyTensorElimination can eliminate tensor.empty inside scf.if blocks. This requires a modeling of equivalence: It is not a per-OpResult property anymore. Instead, it can be specified for each OpOperand and OpResult. This is important because tensor.empty may be eliminated only if all values on the SSA use-def chain to the final consumer (tensor.insert_slice) are equivalent.
• The detection of "returning allocs from a block" can be improved. (Addresses a TODO in assertNoAllocsReturned.) This allows us to bufferize IR such as "yielding a tensor.extract_slice result from an scf.if branch", which currently fails to bufferize because the alloc detection is too conservative.
• Better bufferization of loops. Aliases of the iter_arg can be yielded (even if they are not equivalent) without having to realloc and copy the entire buffer on each iteration.

The above-mentioned examples are not yet implemented with this change. This change just improves the BufferizableOpInterface, its implementations and related helper functions, so that better aliasing information is available for each op.

[aartbik]

LGTM for sparse changes

[springerm]

rebase

[springerm]

rebase

[springerm]

rebase

[springerm]

rebase

[springerm]

rebase

[springerm]

rebase

[nicolasvasilache]

Great generalization, thank you!

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

985

typo

994

Can we add more doc using your debug notation to explain the cases a bit deeper?

E.g.

// aliasingOpOperands(%r) = {(%t, EQUIV, DEFINITE)}
// aliasingOpResults(%t) = {(%r, EQUIV, DEFINITE)}
%r = tensor.insert %f into %t[%idx] : tensor<?xf32>

// aliasingOpOperands(%r) = {(%t0, EQUIV, MAYBE), (%t1, EQUIV, MAYBE)}
// aliasingOpResults(%t0) = {(%r, EQUIV, MAYBE)}
// aliasingOpResults(%t1) = {(%r, EQUIV, MAYBE)}
%r = arith.select %c, %t0, %t1 : tensor<?xf32>

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

 namespace mlir {
 class OpBuilder;
 
 namespace bufferization {
 
 class AnalysisState;
 class BufferizableOpInterface;
 
+/// Specifies a fine-grain relationship between buffers to enable more analysis.
+enum class BufferRelation {
+  None,
+  // TODO: ResultContainsOperand,
+  // TODO: OperandContainsResult,
+  Equivalent
+};
+
+/// Specifies the type of a property.
+enum class Certainty { MayBe, MustBe };
+
+/// An aliasing OpOperand.
+struct AliasingOpOperand {
+  AliasingOpOperand(OpOperand *opOperand, BufferRelation relation)
+      : opOperand(opOperand), relation(relation) {}
+
+  OpOperand *opOperand;
+  BufferRelation relation;
+};
+
+/// An aliasing OpResult.
+struct AliasingOpResult {
+  AliasingOpResult(OpResult opResult, BufferRelation relation)
+      : opResult(opResult), relation(relation) {}
+
+  OpResult opResult;
+  BufferRelation relation;
+};
+
+template <typename T> class AliasList {
+public:
+  /// Create an empty MayBe list of aliases.
+  AliasList<T>() : certainty(Certainty::MayBe) {}
+
+  /// Create an empty list of aliases with the specified certainty.
+  explicit AliasList<T>(Certainty certainty) : certainty(certainty) {}
+
+  /// Create a list of aliases with the specified certainty.
+  AliasList<T>(Certainty certainty, std::initializer_list<T> elems)
+      : certainty(certainty) {
+    for (T alias : elems)
+      addAlias(alias);
+  }
+
+  /// Create a list of aliases with the specified certainty.
+  AliasList<T>(Certainty certainty, SmallVector<T> &&aliases)
+      : aliases(std::move(aliases)), certainty(certainty) {}
+
+  ArrayRef<T> getAliases() const { return aliases; }
+
+  Certainty getCertainty() const { return certainty; }
+
+  size_t getNumAliases() const { return aliases.size(); }
+
+  void addAlias(T alias) { aliases.push_back(alias); }
+
+  auto begin() const { return aliases.begin(); }
+  auto end() const { return aliases.end(); }
+
+private:
+  /// The list of aliases.
+  SmallVector<T> aliases;
+
+  /// Specifies the type of the property.
+  /// * MayBe: Either one of the entries may turn out to alias at runtime. The
+  ///   runtime alias is guaranteed to be in the list.
+  /// * MustBe: Each entry is guaranteed to alias at runtime.
+  const Certainty certainty;
+};
+
+/// A list of aliasing OpOperands.
+using AliasingOpOperandList = AliasList<AliasingOpOperand>;
+/// A list of aliasing OpResults.
+using AliasingOpResultList = AliasList<AliasingOpResult>;
+
 class OpFilter {
 public:
   /// An op filter entry. Filters can be used to specify which ops should be
   /// processed by the bufferization.
   struct Entry {
     /// If the filter function evaluates to `true`, the filter matches.
     using FilterFn = std::function<bool(Operation *)>;
 
struct BufferizationOptions {
   /// Buffer alignment for new memory allocations.
   unsigned int bufferAlignment = 64;
 
   /// Initializer functions for analysis state. These can be used to
   /// initialize dialect-specific analysis state.
   SmallVector<AnalysisStateInitFn> stateInitializers;
 };
 
-/// Specify fine-grain relationship between buffers to enable more analysis.
-enum class BufferRelation {
-  None,
-  // TODO: ResultContainsOperand,
-  // TODO: OperandContainsResult,
-  Equivalent
-};
-
 /// Return `true` if the given value is a BlockArgument of a func::FuncOp.
 bool isFunctionArgument(Value value);
 
 /// AnalysisState provides a variety of helper functions for dealing with
 /// tensor values.
 class AnalysisState {
 public:
-  /// Determine which OpOperand* will alias with `opResult` if the op is
+  /// Determine which OpOperand* will alias with `result` if the op is
   /// bufferized in place. Return all tensor OpOperand* if the op is not
   /// bufferizable.
-  SmallVector<OpOperand *> getAliasingOpOperand(OpResult opResult) const;
+  AliasingOpOperandList getAliasingOpOperands(OpResult result) const;
 
   /// Determine which OpResult will alias with `opOperand` if the op is
   /// bufferized in place. Return all tensor OpResults if the op is not
   /// bufferizable.
-  SmallVector<OpResult> getAliasingOpResult(OpOperand &opOperand) const;
+  AliasingOpResultList getAliasingOpResults(OpOperand &opOperand) const;
 
   /// Return true if `opOperand` bufferizes to a memory read. Return `true` if
   /// the op is not bufferizable.
   bool bufferizesToMemoryRead(OpOperand &opOperand) const;
 
   /// Return true if `opOperand` bufferizes to a memory write. Return true` if
   /// the op is not bufferizable.
   bool bufferizesToMemoryWrite(OpOperand &opOperand) const;
Region *getEnclosingRepetitiveRegion(Value value,
                                      const BufferizationOptions &options);
 
 /// Return the closest enclosing repetitive region around the given block.
 Region *getEnclosingRepetitiveRegion(Block *block,
                                      const BufferizationOptions &options);
 
 namespace detail {
 /// This is the default implementation of
+/// BufferizableOpInterface::getAliasingOpOperands. Should not be called from
+/// other places.
+AliasingOpOperandList defaultGetAliasingOpOperands(OpResult opResult,
+                                                   const AnalysisState &state);
+
+/// This is the default implementation of
 /// BufferizableOpInterface::getBufferType. Should not be called from other
 /// places.
 FailureOr<BaseMemRefType>
 defaultGetBufferType(Value value, const BufferizationOptions &options,
                      const DenseMap<Value, BaseMemRefType> &fixedTypes);
 
 /// This is the default implementation of
 /// BufferizableOpInterface::isRepetitiveRegion. Should not be called from other
 /// places.
 bool defaultIsRepetitiveRegion(BufferizableOpInterface bufferizableOp,
                                unsigned index);
 
-/// This is the default implementation of getAliasingOpOperand in case the
-/// defining op does not implement the BufferizableOpInterface.
-SmallVector<OpOperand *> unknownGetAliasingOpOperand(OpResult opResult);
+/// This is the default implementation of getAliasingOpOperands in case the
+/// owner op does not implement the BufferizableOpInterface.
+AliasingOpOperandList unknownGetAliasingOpOperands(OpResult opResult);
 
-/// This is the default implementation of getAliasingOpResult in case the
+/// This is the default implementation of getAliasingOpResults in case the
 /// defining op does not implement the BufferizableOpInterface.
-SmallVector<OpResult> unknownGetAliasingOpResult(OpOperand &opOperand);
+AliasingOpResultList unknownGetAliasingOpResults(OpOperand &opOperand);
 } // namespace detail
 
 } // namespace bufferization
 } // namespace mlir
 
 MLIR_DECLARE_EXPLICIT_TYPE_ID(mlir::bufferization::AnalysisState)
 
 //===----------------------------------------------------------------------===//
 // Bufferization Interfaces
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h.inc"
 
 #endif // MLIR_DIALECT_BUFFERIZATION_IR_BUFFERIZABLEOPINTERFACE_H_

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

let methods = [
           /*retType=*/"bool",
           /*methodName=*/"isMemoryWrite",
           /*args=*/(ins "::mlir::OpResult":$opResult,
                         "const ::mlir::bufferization::AnalysisState &":$state),
           /*methodBody=*/"",
           /*defaultImplementation=*/[{
             auto bufferizableOp =
                 cast<BufferizableOpInterface>($_op.getOperation());
-            SmallVector<OpOperand*> opOperands =
-              bufferizableOp.getAliasingOpOperand(opResult, state);
-            if (opOperands.empty())
+            AliasingOpOperandList aliases =
+              bufferizableOp.getAliasingOpOperands(opResult, state);
+            if (aliases.getNumAliases() == 0)
               return true;
             return llvm::any_of(
-                opOperands,
-                [&](OpOperand *operand) {
-                  return bufferizableOp.bufferizesToMemoryWrite(*operand,
-                                                                state);
+                aliases,
+                [&](AliasingOpOperand alias) {
+                  return bufferizableOp.bufferizesToMemoryWrite(
+                      *alias.opOperand, state);
                 });
           }]
       >,
       InterfaceMethod<
         /*desc=*/[{
           Return `true` if the given OpOperand must bufferize in-place. Alias
           sets and inplace attributes will be set up accordingly before making
           any other bufferization decisions. This method will never be called on
           OpOperands that do not have a tensor type.
         }],
         /*retType=*/"bool",
         /*methodName=*/"mustBufferizeInPlace",
         /*args=*/(ins "::mlir::OpOperand &":$opOperand,
                       "const ::mlir::bufferization::AnalysisState &":$state),
         /*methodBody=*/"",
         /*defaultImplementation=*/[{
           return false;
         }]
       >,
       InterfaceMethod<
         /*desc=*/[{
-          Return the OpResult that aliases with a given OpOperand when
+          Return the OpResults that may alias with a given OpOperand when
           bufferized in-place. This method will never be called on OpOperands
           that do not have a tensor type.
 
-          Note: This method can return multiple OpResults, indicating that a
-          given OpOperand may at runtime alias with any (or multiple) of the
-          returned OpResults.
+          This method can return multiple OpResults, indicating that a given
+          OpOperand may at runtime alias with any (or multiple) of the returned
+          OpResults.
+
+          This property is specified with a degree of certainty:
+
+          * Certainty::MayBe: The returned set of OpResults is an exhaustive
+            set of possible aliases. At runtime, buffer(opOperand) may alias
+            with any (or none) of the buffers of the returned OpResults. False
+            positives are allowed. On the contrary, omitting potential aliases
+            is incorrect.
+          * Certainty::MustBe: The buffers of returned set of OpResults are
+            guaranteed to alias with buffer(opOperand). False positives are not
+            allowed. Omitting potential aliases is also incorrect.
+
+          Certainty::MustBe is a stronger property than Certainty::MayBe and can
+          enable more efficient bufferization. When possible, aliasing should be
+          specified as a MustBe relationship.
+
+          For each alias, a BufferRelation can be specified:
+
+          * BufferRelation::Equivalent: Both aliases are the exact same buffer.
+            I.e., same size, no offset, same strides.
+          * BufferRelation::None: There is no further information apart from the
+            fact that both buffers alias.
+
+          One possible (conservative) implementation of this interface method
+          that is always safe is to return all tensor OpResults with
+          BufferRelation::None and Certainty::Maybe.
+
+          Examples:
+
+          ```
+          // aliasingOpResults(%t) = MustBe {Equivalent %r}
+          %r = tensor.insert_slice %f into %t : tensor<10xf32>
+
+          // aliasingOpResults(%t) = MustBe {None %r}
+          %r = tensor.extract_slice %t[0]][5][1]
+              : tensor<10xf32> to tensor<5xf32>
+
+          // aliasingOpResults(%t1) = MayBe {Equivalent %r}
+          // aliasingOpResults(%t2) = MayBe {Equivalent %r}
+          %r = arith.select %c, %t1, %t2 : tensor<10xf32>
+
+          // A hypothetical op that bufferizes to roll a dice and based on the
+          // result ot either return buffer(%t) or a newly allocated copy
+          // thereof.
+          // aliasingOpResults(%t1) = MayBe {Equivalent %r}
+          %r = "dummy.alias_or_copy(%t) : (tensor<10xf32>) -> (tensor<10xf32>)"
+          ```
         }],
-        /*retType=*/"::mlir::SmallVector<::mlir::OpResult>",
-        /*methodName=*/"getAliasingOpResult",
+        /*retType=*/"::mlir::bufferization::AliasingOpResultList",
+        /*methodName=*/"getAliasingOpResults",
         /*args=*/(ins "::mlir::OpOperand &":$opOperand,
                       "const ::mlir::bufferization::AnalysisState &":$state),
         /*methodBody=*/"",
         /*defaultImplementation=*/[{
           // Does not have to be implemented for ops without tensor OpOperands.
-          llvm_unreachable("getAliasingOpResult not implemented");
+          assert(opOperand.get().getType().isa<TensorType>() &&
+                 "expected OpOperand with tensor type");
+          llvm_unreachable("getAliasingOpResults not implemented");
         }]
       >,
       InterfaceMethod<
         /*desc=*/[{
           Return the OpOperands that alias with a given OpResult when
           bufferized in-place. This method will never be called on OpResults
           that do not have a tensor type.
 
-          By default, this method is the inverse of `getAliasingOpResult`. Ops
+          By default, this method is the inverse of `getAliasingOpResults`. Ops
           with a region that yield values may want to override this method to
           return the OpOperands that are yielded by the terminator.
 
-          Note: This method can return multiple OpOperands, indicating that the
-          given OpResult may at runtime alias with any (or multiple) of the
-          returned OpOperands. This can be useful for branches and for ops such
-          as `arith.select`.
+          Examples:
+
+          ```
+          // aliasingOpOperands(%r) = MustBe {Equivalent %t}
+          %r = tensor.insert_slice %f into %t : tensor<10xf32>
+
+          // aliasingOpOperands(%r) = MustBe {None %t}
+          %r = tensor.extract_slice %t[0]][5][1]
+              : tensor<10xf32> to tensor<5xf32>
+
+          // aliasingOpOperands(%r) = MayBe {Equivalent %t1, Equivalent %t2}
+          %r = arith.select %c, %t1, %t2 : tensor<10xf32>
+
+          // aliasingOpOperands(%r) = MayBe {}
+          %r = tensor.empty() : tensor<10xf32>
+          ```
         }],
-        /*retType=*/"::mlir::SmallVector<::mlir::OpOperand *>",
-        /*methodName=*/"getAliasingOpOperand",
+        /*retType=*/"::mlir::bufferization::AliasingOpOperandList",
+        /*methodName=*/"getAliasingOpOperands",
         /*args=*/(ins "::mlir::OpResult":$opResult,
                       "const ::mlir::bufferization::AnalysisState &":$state),
         /*methodBody=*/"",
         /*defaultImplementation=*/[{
           assert(opResult.getType().isa<TensorType>() &&
                  "expected OpResult with tensor type");
-          SmallVector<OpOperand *> result;
-          auto bufferizableOp =
-              cast<BufferizableOpInterface>($_op.getOperation());
-          for (OpOperand &opOperand : $_op.getOperation()->getOpOperands()) {
-            if (!opOperand.get().getType().isa<TensorType>())
-              continue;
-            SmallVector<OpResult> aliasingOpResults =
-                bufferizableOp.getAliasingOpResult(opOperand, state);
-            if (llvm::is_contained(aliasingOpResults, opResult))
-              result.push_back(&opOperand);
-          }
-          return result;
-        }]
-      >,
-      InterfaceMethod<
-        /*desc=*/[{
-          Return the buffer relation between the given OpResult and its aliasing
-          OpOperands when bufferized in-place. Most OpOperands have an
-          "equivalence" relation. This method will never be called on OpResults
-          that do not have a tensor type. It will also never be called on
-          OpResults that do not have at least one aliasing OpOperand.
-
-          TODO: Support other relations such as "OpOperand is included in
-          OpResult".
-        }],
-        /*retType=*/"::mlir::bufferization::BufferRelation",
-        /*methodName=*/"bufferRelation",
-        /*args=*/(ins "::mlir::OpResult":$opResult,
-                      "const ::mlir::bufferization::AnalysisState &":$state),
-        /*methodBody=*/"",
-        /*defaultImplementation=*/[{
-          // Does not have to be implemented for ops without tensor OpResults
-          // that have an aliasing OpOperand.
-          llvm_unreachable("bufferRelation not implemented");
+          return ::mlir::bufferization::detail::defaultGetAliasingOpOperands(
+              opResult, state);
         }]
       >,
       InterfaceMethod<
         /*desc=*/[{
           Resolve all inplacability conflicts by inserting explicit
           `bufferization.alloc_tensor` ops. Examples of inplacability conflicts
           are read-after-write conflicts or writes into non-writable buffers.
 
let methods = [
           Bufferize this op, i.e., rewrite it into a memref-based equivalent.
           Buffers of tensor SSA values can be retrieved via `getBuffer`.
           Uses of tensor results of the existing tensor op can be replaced with
           `replaceOpWithBufferizedValues` or `replaceOpWithNewBufferizedOp`.
           These two functions automatically handle the tensor-to-memref type
           conversion.
 
           The implementation of this method must be consistent with the
-          remaining methods, in particular `getAliasingOpOperand`. I.e., a
+          remaining methods, in particular `getAliasingOpOperands`. I.e., a
           tensor result `r` may only be replaced with:
-          a) A buffer that aliases one of buffers in getAliasingOpOperand(r).
+
+          a) A buffer that aliases one of buffers in getAliasingOpOperands(r).
           b) Or: A newly allocated buffer.
 
+          If `getAliasingOpOperands` is specified as Certainty::MustBe, the
+          second option is not available.
+
           This method will never be called on ops that do not have at least one
           tensor operand/result.
 
           The return value of this method indicates whether there was an error
           while bufferizing this op (such as failing to create a new buffer
           allocation op). The bufferization driver immediately stops bufferizing
           the input IR and returns `failure` in that case. If this op is
           expected to survive bufferization, `success` should be returned
let extraClassDeclaration = [{
     ///
     /// Examples of such ops are `tensor.extract_slice` and `tensor.cast`.
     bool bufferizesToAliasOnly(OpOperand &opOperand,
                                const AnalysisState &state) {
       auto bufferizableOp =
           cast<BufferizableOpInterface>(getOperation());
       return !bufferizableOp.bufferizesToMemoryRead(opOperand, state)
           && !bufferizableOp.bufferizesToMemoryWrite(opOperand, state)
-          && !bufferizableOp.getAliasingOpResult(opOperand, state).empty();
+          && bufferizableOp.getAliasingOpResults(opOperand, state)
+              .getNumAliases() != 0;
     }
   }];
 }
 
 #endif  // BUFFERIZABLE_OP_INTERFACE

mlir/include/mlir/Dialect/Bufferization/IR/BufferizationOps.td

let extraClassDeclaration = [{
     bool bufferizesToAllocation(OpResult opResult) { return true; }
 
     bool bufferizesToMemoryRead(OpOperand &opOperand,
                                 const AnalysisState &state);
 
     bool bufferizesToMemoryWrite(OpOperand &opOperand,
                                  const AnalysisState &state);
 
-    SmallVector<OpResult> getAliasingOpResult(
+    AliasingOpResultList getAliasingOpResults(
         OpOperand &opOperand, const AnalysisState &state);
 
     FailureOr<BaseMemRefType> getBufferType(
         Value value, const BufferizationOptions &options,
         const DenseMap<Value, BaseMemRefType> &fixedTypes);
 
     RankedTensorType getType() {
       return getResult().getType().cast<RankedTensorType>();
bool bufferizesToMemoryRead(OpOperand &opOperand,
       return false;
     }
 
     bool bufferizesToMemoryWrite(OpOperand &opOperand,
                                  const AnalysisState &state) const {
       return false;
     }
 
-    SmallVector<OpResult> getAliasingOpResult(
+    AliasingOpResultList getAliasingOpResults(
         OpOperand &opOperand, const AnalysisState &state) const {
       return {};
     }
 
     LogicalResult bufferize(RewriterBase &rewriter,
                             const BufferizationOptions &options);
   }];
 }
let extraClassDeclaration = [{
     }
 
     bool mustBufferizeInPlace(OpOperand &opOperand,
                               const AnalysisState &state) const {
       // ToMemrefOps always bufferize inplace.
       return true;
     }
 
-    SmallVector<OpResult> getAliasingOpResult(
+    AliasingOpResultList getAliasingOpResults(
         OpOperand &opOperand, const AnalysisState &state) const {
       return {};
     }
 
     LogicalResult bufferize(RewriterBase &rewriter,
                             const BufferizationOptions &options);
   }];
 
   let assemblyFormat = "$tensor attr-dict `:` type($memref)";
 
   let hasFolder = 1;
   let hasCanonicalizer = 1;
 }
 
 #endif // BUFFERIZATION_OPS

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

struct DstBufferizableOpInterfaceExternalModel
 
   bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
                                const AnalysisState &state) const {
     // Only outputs bufferize to a memory write.
     auto dstOp = cast<DestinationStyleOpInterface>(op);
     return dstOp.isDpsInit(&opOperand);
   }
 
-  SmallVector<OpResult> getAliasingOpResult(Operation *op, OpOperand &opOperand,
+  AliasingOpResultList getAliasingOpResults(Operation *op, OpOperand &opOperand,
                                             const AnalysisState &state) const {
     // Output operands alias with their respective tied OpResults.
     auto dstOp = cast<DestinationStyleOpInterface>(op);
     if (dstOp.isDpsInit(&opOperand))
-      return {dstOp.getTiedOpResult(&opOperand)};
+      return {
+          Certainty::MustBe,
+          {{dstOp.getTiedOpResult(&opOperand), BufferRelation::Equivalent}}};
     return {};
   }
-
-  BufferRelation bufferRelation(Operation *op, OpResult opResult,
-                                const AnalysisState &state) const {
-    assert(isa<DestinationStyleOpInterface>(op) &&
-           "expected that op implements DestinationStyleOpInterface");
-    return BufferRelation::Equivalent;
-  }
 };
 
 } // namespace bufferization
 } // namespace mlir
 
 #endif // MLIR_DIALECT_BUFFERIZATION_IR_DSTBUFFERIZABLEOPINTERFACEIMPL_H_

mlir/lib/Dialect/Arith/Transforms/BufferizableOpInterfaceImpl.cpp

bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
     return false;
   }
 
   bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
                                const AnalysisState &state) const {
     return false;
   }
 
-  SmallVector<OpResult> getAliasingOpResult(Operation *op, OpOperand &opOperand,
+  AliasingOpResultList getAliasingOpResults(Operation *op, OpOperand &opOperand,
                                             const AnalysisState &state) const {
-    return {op->getResult(0)};
-  }
-
-  BufferRelation bufferRelation(Operation *op, OpResult opResult,
-                                const AnalysisState &state) const {
-    return BufferRelation::Equivalent;
+    return {Certainty::MustBe,
+            {{op->getResult(0), BufferRelation::Equivalent}}};
   }
 
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
                           const BufferizationOptions &options) const {
     auto castOp = cast<arith::IndexCastOp>(op);
     auto resultTensorType = castOp.getType().cast<TensorType>();
 
     FailureOr<Value> source = getBuffer(rewriter, castOp.getIn(), options);
bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
     return false;
   }
 
   bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
                                const AnalysisState &state) const {
     return false;
   }
 
-  SmallVector<OpResult> getAliasingOpResult(Operation *op, OpOperand &opOperand,
+  AliasingOpResultList getAliasingOpResults(Operation *op, OpOperand &opOperand,
                                             const AnalysisState &state) const {
-    return {op->getOpResult(0) /*result*/};
-  }
-
-  SmallVector<OpOperand *>
-  getAliasingOpOperand(Operation *op, OpResult opResult,
-                       const AnalysisState &state) const {
-    return {&op->getOpOperand(1) /*true_value*/,
-            &op->getOpOperand(2) /*false_value*/};
+    return {Certainty::MayBe,
+            {{op->getOpResult(0) /*result*/, BufferRelation::Equivalent}}};
   }
 
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
                           const BufferizationOptions &options) const {
     auto selectOp = cast<arith::SelectOp>(op);
     Location loc = selectOp.getLoc();
 
     // TODO: It would be more efficient to copy the result of the `select` op
getBufferType(Operation *op, Value value, const BufferizationOptions &options,
     // If the buffers have different types, they differ only in their layout
     // map.
     auto memrefType = trueType->cast<MemRefType>();
     return getMemRefTypeWithFullyDynamicLayout(
         RankedTensorType::get(memrefType.getShape(),
                               memrefType.getElementType()),
         memrefType.getMemorySpace());
   }
-
-  BufferRelation bufferRelation(Operation *op, OpResult opResult,
-                                const AnalysisState &state) const {
-    return BufferRelation::None;
-  }
 };
 
 } // namespace
 
 void mlir::arith::registerBufferizableOpInterfaceExternalModels(
     DialectRegistry &registry) {
   registry.addExtension(+[](MLIRContext *ctx, ArithDialect *dialect) {
     ConstantOp::attachInterface<ConstantOpInterface>(*ctx);
     IndexCastOp::attachInterface<IndexCastOpInterface>(*ctx);
     SelectOp::attachInterface<SelectOpInterface>(*ctx);
   });
 }

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

for (OpOperand &opOperand : op->getOpOperands()) {
     Type operandType = opOperand.get().getType();
     if (!operandType.isa<TensorType>())
       continue;
     if (state.isInPlace(opOperand))
       continue;
     if (operandType.isa<UnrankedTensorType>())
       return op->emitError("copies of unranked tensors are not supported");
 
-    SmallVector<OpResult> aliasingOpResults =
-        state.getAliasingOpResult(opOperand);
+    AliasingOpResultList aliasingOpResults =
+        state.getAliasingOpResults(opOperand);
     // Is the result yielded from a block? Or are deallocations turned off
     // entirely? In either case, mark the allocation as "escaping", so that it
     // will not be deallocated.
     bool escape = !state.getOptions().createDeallocs ||
-                  llvm::any_of(aliasingOpResults, [&](Value v) {
-                    return state.isTensorYielded(v);
+                  llvm::any_of(aliasingOpResults, [&](AliasingOpResult a) {
+                    return state.isTensorYielded(a.opResult);
                   });
 
-    if (aliasingOpResults.size() == 1 &&
+    if (aliasingOpResults.getNumAliases() == 1 &&
         !state.bufferizesToMemoryWrite(opOperand) &&
-        state.getAliasingOpOperand(aliasingOpResults.front()).size() == 1) {
+        state.getAliasingOpOperands(aliasingOpResults.getAliases()[0].opResult)
+                .getNumAliases() == 1) {
       // The op itself does not write but may create exactly one alias. Instead
       // of copying the OpOperand, copy the OpResult. The OpResult can sometimes
       // be smaller than the OpOperand (e.g., in the case of an extract_slice,
       // where the result is usually a smaller part of the source).
-      outOfPlaceOpResults.push_back(aliasingOpResults.front());
+      OpResult opResult = aliasingOpResults.getAliases()[0].opResult;
+      outOfPlaceOpResults.push_back(opResult);
       if (!state.canOmitTensorCopy(opOperand))
-        copiedOpResults.insert(aliasingOpResults.front());
+        copiedOpResults.insert(opResult);
       if (escape)
-        escapingOpResultCopies.insert(aliasingOpResults.front());
+        escapingOpResultCopies.insert(opResult);
     } else {
       // In all other cases, make a copy of the OpOperand.
       outOfPlaceOpOperands.push_back(&opOperand);
       if (!state.canOmitTensorCopy(opOperand))
         copiedOpOperands.insert(&opOperand);
       if (escape)
         escapingOpOperandCopies.insert(&opOperand);
     }
static void setInsertionPointAfter(OpBuilder &b, Value value) {
   } else {
     b.setInsertionPointAfter(value.getDefiningOp());
   }
 }
 
 /// Determine which OpOperand* will alias with `opResult` if the op is
 /// bufferized in place. Return all tensor OpOperand* if the op is not
 /// bufferizable.
-SmallVector<OpOperand *>
-AnalysisState::getAliasingOpOperand(OpResult opResult) const {
+AliasingOpOperandList
+AnalysisState::getAliasingOpOperands(OpResult opResult) const {
   if (Operation *op = opResult.getDefiningOp())
     if (auto bufferizableOp = getOptions().dynCastBufferizableOp(op))
-      return bufferizableOp.getAliasingOpOperand(opResult, *this);
+      return bufferizableOp.getAliasingOpOperands(opResult, *this);
 
   // The op is not bufferizable.
-  return detail::unknownGetAliasingOpOperand(opResult);
+  return detail::unknownGetAliasingOpOperands(opResult);
 }
 
 /// Determine which OpResult will alias with `opOperand` if the op is bufferized
 /// in place. Return all tensor OpResults if the op is not bufferizable.
-SmallVector<OpResult>
-AnalysisState::getAliasingOpResult(OpOperand &opOperand) const {
+AliasingOpResultList
+AnalysisState::getAliasingOpResults(OpOperand &opOperand) const {
   if (auto bufferizableOp =
           getOptions().dynCastBufferizableOp(opOperand.getOwner()))
-    return bufferizableOp.getAliasingOpResult(opOperand, *this);
+    return bufferizableOp.getAliasingOpResults(opOperand, *this);
 
   // The op is not bufferizable.
-  return detail::unknownGetAliasingOpResult(opOperand);
+  return detail::unknownGetAliasingOpResults(opOperand);
 }
 
 /// Return true if `opOperand` bufferizes to a memory read. Return `true` if the
 /// op is not bufferizable.
 bool AnalysisState::bufferizesToMemoryRead(OpOperand &opOperand) const {
   if (auto bufferizableOp =
           getOptions().dynCastBufferizableOp(opOperand.getOwner()))
     return bufferizableOp.bufferizesToMemoryRead(opOperand, *this);
bool AnalysisState::isValueRead(Value value) const {
   SmallVector<OpOperand *> workingSet;
   for (OpOperand &use : value.getUses())
     workingSet.push_back(&use);
 
   while (!workingSet.empty()) {
     OpOperand *uMaybeReading = workingSet.pop_back_val();
     // Skip over all ops that neither read nor write (but create an alias).
     if (bufferizesToAliasOnly(*uMaybeReading))
-      for (OpResult opResult : getAliasingOpResult(*uMaybeReading))
-        for (OpOperand &use : opResult.getUses())
+      for (AliasingOpResult alias : getAliasingOpResults(*uMaybeReading))
+        for (OpOperand &use : alias.opResult.getUses())
           workingSet.push_back(&use);
     if (bufferizesToMemoryRead(*uMaybeReading))
       return true;
   }
 
   return false;
 }
 
while (!workingSet.empty()) {
     if (condition(value) || value.isa<BlockArgument>()) {
       result.insert(value);
       continue;
     }
 
     OpResult opResult = value.cast<OpResult>();
     BufferizableOpInterface bufferizableOp =
         options.dynCastBufferizableOp(opResult.getDefiningOp());
-    SmallVector<OpOperand *> opOperands = getAliasingOpOperand(opResult);
+    AliasingOpOperandList aliases = getAliasingOpOperands(opResult);
 
     // 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.
-    // * There is an OpOperand, but it is not an equivalent tensor (only if
-    //   `followEquivalentOnly` is set).
-    if (!bufferizableOp || opOperands.empty() ||
-        (followEquivalentOnly &&
-         bufferizableOp.bufferRelation(opResult, *this) !=
-             BufferRelation::Equivalent)) {
+    if (!bufferizableOp || aliases.getNumAliases() == 0) {
       result.insert(value);
       continue;
     }
 
-    for (OpOperand *o : opOperands)
-      workingSet.insert(o->get());
+    for (AliasingOpOperand a : aliases) {
+      if (followEquivalentOnly && a.relation != BufferRelation::Equivalent) {
+        // Stop iterating if `followEquivalentOnly` is set but the alias is not
+        // equivalent.
+        result.insert(value);
+      } else {
+        workingSet.insert(a.opOperand->get());
+      }
+    }
   }
 
   return result;
 }
 
 // Find the Values of the last preceding write of a given Value.
 llvm::SetVector<Value>
 AnalysisState::findLastPrecedingWrite(Value value) const {
if (hasUndefinedContents(&opOperand))
     return true;
 
   // Do not copy if the buffer of the tensor is entirely overwritten (with
   // values that do not depend on the old tensor).
   if (bufferizesToMemoryWrite(opOperand) && !bufferizesToMemoryRead(opOperand))
     return true;
 
   // Do not copy if the tensor is never read.
-  SmallVector<OpResult> aliasingOpResults = getAliasingOpResult(opOperand);
+  AliasingOpResultList aliases = getAliasingOpResults(opOperand);
   if (!bufferizesToMemoryRead(opOperand) &&
-      llvm::none_of(aliasingOpResults,
-                    [&](OpResult opResult) { return isValueRead(opResult); }))
+      llvm::none_of(
+          aliases, [&](AliasingOpResult a) { return isValueRead(a.opResult); }))
     return true;
 
   // Default: Cannot omit the copy.
   return false;
 }
 
 bool AnalysisState::isInPlace(OpOperand &opOperand) const {
   // ToMemrefOps are always in-place.
while (!worklist.empty()) {
     // Check if the op is returning/yielding.
     if (isRegionReturnLike(op))
       return true;
 
     // Add all aliasing OpResults to the worklist.
     // Note: In the absence of detailed analysis information (e.g., there may be
     // no function call analysis information), this `getAliasingOpResult` is
     // conservative and may report additional OpResults as potentially aliasing.
-    for (OpResult opResult : getAliasingOpResult(*operand))
-      for (OpOperand &use : opResult.getUses())
+    for (AliasingOpResult alias : getAliasingOpResults(*operand))
+      for (OpOperand &use : alias.opResult.getUses())
         worklist.push_back(&use);
   }
 
   // No ReturnLike op found: The value is not yielded.
   return false;
 }
 
 // bufferization.to_memref is not allowed to change the rank.
return MemRefType::get(rankedTensorType.getShape(),
                          rankedTensorType.getElementType(), layout,
                          memorySpace);
 }
 
 //===----------------------------------------------------------------------===//
 // Default implementations of interface methods
 //===----------------------------------------------------------------------===//
 
+// Compute the AliasingOpOperandList for a given OpResult based on
+// getAliasingOpResults.
+AliasingOpOperandList bufferization::detail::defaultGetAliasingOpOperands(
+    OpResult opResult, const AnalysisState &state) {
+  Operation *op = opResult.getDefiningOp();
+  SmallVector<AliasingOpOperand> result;
+  std::optional<Certainty> certainty = {};
+  for (OpOperand &opOperand : op->getOpOperands()) {
+    if (!opOperand.get().getType().isa<TensorType>())
+      continue;
+    AliasingOpResultList aliasingOpResults =
+        state.getAliasingOpResults(opOperand);
+    for (const auto &it : aliasingOpResults) {
+      if (it.opResult == opResult) {
+        result.emplace_back(&opOperand, it.relation);
+        if (!certainty.has_value())
+          certainty = aliasingOpResults.getCertainty();
+      }
+      break;
+    }
+  }
+  return AliasingOpOperandList(
+      certainty.has_value() ? *certainty : Certainty::MayBe, std::move(result));
+}
+
 FailureOr<BaseMemRefType> bufferization::detail::defaultGetBufferType(
     Value value, const BufferizationOptions &options,
     const DenseMap<Value, BaseMemRefType> &fixedTypes) {
   assert(value.getType().isa<TensorType>() && "expected tensor type");
 
   // No further analysis is possible for a block argument.
   if (value.isa<BlockArgument>())
     return bufferization::getMemRefType(value, options);
 
   // Value is an OpResult.
   Operation *op = getOwnerOfValue(value);
   auto opResult = value.cast<OpResult>();
-  auto bufferizableOp = cast<BufferizableOpInterface>(op);
   AnalysisState state(options);
-  auto aliasingOperands = bufferizableOp.getAliasingOpOperand(opResult, state);
-  if (!aliasingOperands.empty() &&
-      bufferizableOp.bufferRelation(opResult, state) ==
-          BufferRelation::Equivalent) {
+  AliasingOpOperandList aliases = state.getAliasingOpOperands(opResult);
+  if (aliases.getNumAliases() > 0 &&
+      aliases.getAliases()[0].relation == BufferRelation::Equivalent) {
     // If the OpResult has an equivalent OpOperand, both OpResult and
     // OpOperand bufferize to the exact same buffer type.
-    Value equivalentOperand = aliasingOperands.front()->get();
+    Value equivalentOperand = aliases.getAliases().front().opOperand->get();
     return getBufferType(equivalentOperand, options, fixedTypes);
   }
 
   // If we do not know the memory space and there is no default memory space,
   // report a failure.
   if (!options.defaultMemorySpace.has_value())
     return op->emitError("could not infer memory space");
 
   return getMemRefType(value, options, /*layout=*/{},
                        *options.defaultMemorySpace);
 }
 
 bool bufferization::detail::defaultIsRepetitiveRegion(
     BufferizableOpInterface bufferizableOp, unsigned index) {
   assert(index < bufferizableOp->getNumRegions() && "invalid region index");
   auto regionInterface =
       dyn_cast<RegionBranchOpInterface>(bufferizableOp.getOperation());
   if (!regionInterface)
     return false;
   return regionInterface.isRepetitiveRegion(index);
 }
 
-SmallVector<OpOperand *>
-bufferization::detail::unknownGetAliasingOpOperand(OpResult opResult) {
+AliasingOpOperandList
+bufferization::detail::unknownGetAliasingOpOperands(OpResult opResult) {
   // Conservatively assume that everything is aliasing.
-  SmallVector<OpOperand *> r;
+  AliasingOpOperandList r(Certainty::MayBe);
   for (OpOperand &operand : opResult.getDefiningOp()->getOpOperands())
     if (operand.get().getType().isa<TensorType>())
-      r.push_back(&operand);
+      r.addAlias({&operand, BufferRelation::None});
   return r;
 }
 
-SmallVector<OpResult>
-bufferization::detail::unknownGetAliasingOpResult(OpOperand &opOperand) {
+AliasingOpResultList
+bufferization::detail::unknownGetAliasingOpResults(OpOperand &opOperand) {
   // Conservatively assume that everything is aliasing.
-  SmallVector<OpResult> r;
+  AliasingOpResultList r(Certainty::MayBe);
   for (OpResult result : opOperand.getOwner()->getOpResults())
     if (result.getType().isa<TensorType>())
-      r.push_back(result);
+      r.addAlias({result, BufferRelation::None});
   return r;
 }

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

 
 bool AllocTensorOp::bufferizesToMemoryWrite(OpOperand &opOperand,
                                             const AnalysisState &state) {
   assert(opOperand.getOperandNumber() == getNumOperands() - 1 &&
          "expected copy operand");
   return false;
 }
 
-SmallVector<OpResult>
-AllocTensorOp::getAliasingOpResult(OpOperand &opOperand,
+AliasingOpResultList
+AllocTensorOp::getAliasingOpResults(OpOperand &opOperand,
                                     const AnalysisState &state) {
   // This is a new allocation. It does not alias with any other buffer.
   return {};
 }
 
 FailureOr<BaseMemRefType> AllocTensorOp::getBufferType(
     Value value, const BufferizationOptions &options,
     const DenseMap<Value, BaseMemRefType> &fixedTypes) {
   assert(value == getResult() && "invalid value");

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

if (getFuncOpAnalysisState(state, funcOp) != FuncOpAnalysisState::Analyzed)
       // FuncOp not analyzed yet. Assume that OpOperand is written.
       return true;
 
     const FuncAnalysisState &funcState = getFuncAnalysisState(state);
     return funcState.writtenBbArgs.lookup(funcOp).contains(
         opOperand.getOperandNumber());
   }
 
-  SmallVector<OpResult> getAliasingOpResult(Operation *op, OpOperand &opOperand,
+  AliasingOpResultList getAliasingOpResults(Operation *op, OpOperand &opOperand,
                                             const AnalysisState &state) const {
     func::CallOp callOp = cast<func::CallOp>(op);
     FuncOp funcOp = getCalledFunction(callOp);
     assert(funcOp && "expected CallOp to a FuncOp");
     if (getFuncOpAnalysisState(state, funcOp) != FuncOpAnalysisState::Analyzed)
       // FuncOp not analyzed yet. Any OpResult may be aliasing.
-      return detail::unknownGetAliasingOpResult(opOperand);
+      return detail::unknownGetAliasingOpResults(opOperand);
 
     // Get aliasing results from state.
     const FuncAnalysisState &funcState = getFuncAnalysisState(state);
     auto aliasingReturnVals =
         funcState.aliasingReturnVals.lookup(funcOp).lookup(
             opOperand.getOperandNumber());
-    SmallVector<OpResult> result;
-    for (int64_t resultIdx : aliasingReturnVals)
-      result.push_back(callOp->getOpResult(resultIdx));
-    return result;
-  }
 
-  BufferRelation bufferRelation(Operation *op, OpResult opResult,
-                                const AnalysisState &state) const {
-    func::CallOp callOp = cast<func::CallOp>(op);
-    FuncOp funcOp = getCalledFunction(callOp);
-    assert(funcOp && "expected CallOp to a FuncOp");
-    if (getFuncOpAnalysisState(state, funcOp) !=
-        FuncOpAnalysisState::Analyzed) {
-      // Function not analyzed yet. The conservative answer is "None".
-      return BufferRelation::None;
-    }
-
-    const FuncAnalysisState &funcState = getFuncAnalysisState(state);
-    std::optional<int64_t> maybeEquiv =
-        getEquivalentFuncArgIdx(funcOp, funcState, opResult.getResultNumber());
-    if (maybeEquiv) {
-#ifndef NDEBUG
-      SmallVector<OpOperand *> aliasingOpOperands =
-          getAliasingOpOperand(op, opResult, state);
-      assert(aliasingOpOperands.size() == 1 &&
-             "expected exactly 1 aliasing OpOperand");
-      assert(aliasingOpOperands.front()->getOperandNumber() == *maybeEquiv &&
+    // Check if the aliasing OpResult is equivalent to the OpOperand.
+    std::optional<int64_t> equivalent = {};
+    if (aliasingReturnVals.size() == 1) {
+      equivalent = getEquivalentFuncArgIdx(funcOp, funcState,
+                                           aliasingReturnVals.front());
+      assert((!equivalent.has_value() ||
+              *equivalent == opOperand.getOperandNumber()) &&
              "inconsistent analysis state");
-#endif
-      return BufferRelation::Equivalent;
     }
-    return BufferRelation::None;
+    AliasingOpResultList result(equivalent.has_value() ? Certainty::MustBe
+                                                       : Certainty::MayBe);
+    for (int64_t resultIdx : aliasingReturnVals)
+      result.addAlias({callOp->getOpResult(resultIdx),
+                       equivalent.has_value() ? BufferRelation::Equivalent
+                                              : BufferRelation::None});
+    return result;
   }
 
   /// All function arguments are writable. It is the responsibility of the
   /// CallOp to insert buffer copies where necessary.
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
                           const BufferizationOptions &options) const {
     func::CallOp callOp = cast<func::CallOp>(op);
     unsigned numResults = callOp.getNumResults();
bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
     return true;
   }
 
   bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
                                const AnalysisState &state) const {
     return false;
   }
 
-  SmallVector<OpResult> getAliasingOpResult(Operation *op, OpOperand &opOperand,
+  AliasingOpResultList getAliasingOpResults(Operation *op, OpOperand &opOperand,
                                             const AnalysisState &state) const {
     return {};
   }
 
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
                           const BufferizationOptions &options) const {
 #ifndef NDEBUG
     auto returnOp = cast<func::ReturnOp>(op);

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

 }
 
 /// Set the inPlace bufferization spec to true.
 void BufferizationAliasInfo::bufferizeInPlace(OpOperand &operand,
                                               AnalysisState &state) {
   if (inplaceBufferized.contains(&operand))
     return;
   markInPlace(operand);
-  for (OpResult result : state.getAliasingOpResult(operand))
-    aliasInfo.unionSets(result, operand.get());
+  for (AliasingOpResult alias : state.getAliasingOpResults(operand))
+    aliasInfo.unionSets(alias.opResult, operand.get());
   ++statNumTensorInPlace;
 }
 
 /// Set the inPlace bufferization spec to false.
 void BufferizationAliasInfo::bufferizeOutOfPlace(OpOperand &operand) {
   assert(!inplaceBufferized.contains(&operand) &&
          "OpOperand was already decided to bufferize inplace");
   ++statNumTensorOutOfPlace;
for (OpOperand *uConflictingWrite : usesWrite) {
             // conflictingWritingOp happens outside of the block. No
             // conflict.
             continue;
           }
         }
 
         // No conflict if the conflicting write and the last write are the same
         // use.
-        SmallVector<OpResult> aliasingOpResult =
-            state.getAliasingOpResult(*uConflictingWrite);
-        if (aliasingOpResult.size() == 1 && aliasingOpResult[0] == lastWrite) {
+        AliasingOpResultList aliases =
+            state.getAliasingOpResults(*uConflictingWrite);
+        if (aliases.getNumAliases() == 1 &&
+            aliases.getAliases()[0].opResult == lastWrite) {
           LLVM_DEBUG(llvm::dbgs()
                      << "    no conflict: last write and write are same\n");
           continue;
         }
 
         // All requirements are met. Conflict found!
 
         if (options.printConflicts)
for (auto &use : alias.getUses()) {
       // but its aliasing result is eventually fed into an op that does.
       //
       // Note: This is considered a "read" only if the use does not bufferize to
       // a memory write. (We already ruled out memory reads. In case of a memory
       // write, the buffer would be entirely overwritten; in the above example
       // there would then be no flow of data from the extract_slice operand to
       // its result's uses.)
       if (!state.bufferizesToMemoryWrite(use)) {
-        SmallVector<OpResult> opResults = state.getAliasingOpResult(use);
-        if (llvm::any_of(opResults,
-                         [&](OpResult r) { return state.isValueRead(r); }))
+        AliasingOpResultList aliases = state.getAliasingOpResults(use);
+        if (llvm::any_of(aliases, [&](AliasingOpResult a) {
+              return state.isValueRead(a.opResult);
+            }))
           res.insert(&use);
       }
     }
   });
 }
 
 /// Return true if bufferizing `operand` inplace would create a conflict. A read
 /// R and a write W of the same alias set is a conflict if inplace bufferization

 static bool wouldCreateReadAfterWriteInterference(
     OpOperand &operand, const DominanceInfo &domInfo, AnalysisState &state,
     const BufferizationAliasInfo &aliasInfo,
     bool checkConsistencyOnly = false) {
   // Collect reads and writes of all aliases of OpOperand and OpResult.
   DenseSet<OpOperand *> usesRead, usesWrite;
   getAliasingReads(usesRead, operand.get(), aliasInfo, state);
   getAliasingInplaceWrites(usesWrite, operand.get(), aliasInfo, state);
-  for (OpResult result : state.getAliasingOpResult(operand)) {
-    getAliasingReads(usesRead, result, aliasInfo, state);
-    getAliasingInplaceWrites(usesWrite, result, aliasInfo, state);
+  for (AliasingOpResult alias : state.getAliasingOpResults(operand)) {
+    getAliasingReads(usesRead, alias.opResult, aliasInfo, state);
+    getAliasingInplaceWrites(usesWrite, alias.opResult, aliasInfo, state);
   }
   if (!checkConsistencyOnly && state.bufferizesToMemoryWrite(operand))
     usesWrite.insert(&operand);
 
   return hasReadAfterWriteInterference(usesRead, usesWrite, domInfo, state,
                                        aliasInfo);
 }
 
while (!worklist.empty()) {
     }
 
     // If `nextVal` is not a BlockArgument: End of use-def chain reached.
     auto opResult = nextVal.dyn_cast<OpResult>();
     if (!opResult)
       continue;
 
     // Follow reverse SSA use-def chain.
-    SmallVector<OpOperand *> aliasingOpOperands =
-        state.getAliasingOpOperand(opResult);
-    for (OpOperand *opOperand : aliasingOpOperands)
-      if (aliasInfo.isInPlace(*opOperand) || currentOpOperand == opOperand)
-        worklist.push_back(opOperand->get());
+    for (AliasingOpOperand alias : state.getAliasingOpOperands(opResult))
+      if (aliasInfo.isInPlace(*alias.opOperand) ||
+          currentOpOperand == alias.opOperand)
+        worklist.push_back(alias.opOperand->get());
   }
   return false;
 }
 
 /// Return true if bufferizing `operand` inplace would create a write to a
 /// non-writable buffer.
 static bool wouldCreateWriteToNonWritableBuffer(
     OpOperand &operand, const BufferizationAliasInfo &aliasInfo,
     OneShotAnalysisState &state, bool checkConsistencyOnly = false) {
   // Collect writes of all aliases of OpOperand and OpResult.
   DenseSet<OpOperand *> usesWrite;
   getAliasingInplaceWrites(usesWrite, operand.get(), aliasInfo, state);
-  for (OpResult result : state.getAliasingOpResult(operand)) {
-    getAliasingInplaceWrites(usesWrite, result, aliasInfo, state);
+  for (AliasingOpResult alias : state.getAliasingOpResults(operand)) {
+    getAliasingInplaceWrites(usesWrite, alias.opResult, aliasInfo, state);
   }
   if (!checkConsistencyOnly && state.bufferizesToMemoryWrite(operand))
     usesWrite.insert(&operand);
 
   // Assuming that `operand` bufferizes in-place: For each write (to each
   // alias), check if there is a non-writable tensor in the reverse SSA use-def
   // chain.
   for (OpOperand *uWrite : usesWrite) {
static LogicalResult inPlaceAnalysis(Operation *op,
 
   return inPlaceAnalysis(ops, aliasInfo, state, domInfo, analysisFuzzerSeed);
 }
 
 /// Analyze equivalence of tied OpResult/OpOperand pairs of the given ops.
 static void equivalenceAnalysis(SmallVector<Operation *> &ops,
                                 BufferizationAliasInfo &aliasInfo,
                                 AnalysisState &state) {
-  for (Operation *op : ops)
-    if (auto bufferizableOp = state.getOptions().dynCastBufferizableOp(op))
-      for (OpResult opResult : op->getOpResults())
-        if (opResult.getType().isa<TensorType>())
-          for (OpOperand *opOperand :
-               bufferizableOp.getAliasingOpOperand(opResult, state))
-            if (state.isInPlace(*opOperand))
-              if (bufferizableOp.bufferRelation(opResult, state) ==
-                  BufferRelation::Equivalent)
-                aliasInfo.unionEquivalenceClasses(opResult, opOperand->get());
+  for (Operation *op : ops) {
+    for (OpOperand &opOperand : op->getOpOperands()) {
+      if (opOperand.get().getType().isa<TensorType>()) {
+        if (!state.isInPlace(opOperand))
+          // Out-of-place OpOperands bufferize to new allocations and do not
+          // union equivalence sets.
+          continue;
+        AliasingOpResultList aliases = state.getAliasingOpResults(opOperand);
+        if (aliases.getCertainty() == Certainty::MayBe)
+          // If it is not certain that the alias will materialize at runtime,
+          // we have to be conservative.
+          continue;
+        for (AliasingOpResult alias : aliases) {
+          if (alias.relation == BufferRelation::Equivalent)
+            aliasInfo.unionEquivalenceClasses(alias.opResult, opOperand.get());
+        }
+      }
+    }
+  }
 }
 
 /// Analyze equivalence of tied OpResult/OpOperand pairs of all ops contained
 /// in `op`.
 static void equivalenceAnalysis(Operation *op,
                                 BufferizationAliasInfo &aliasInfo,
                                 AnalysisState &state) {
   // Traverse ops in PostOrder: Nested ops first, then enclosing ops.
   SmallVector<Operation *> ops;
   op->walk<WalkOrder::PostOrder>([&](Operation *op) {
     // No tensors => no buffers.
-    if (none_of(op->getResultTypes(), isaTensor))
+    if (none_of(op->getOperandTypes(), isaTensor))
       return;
     ops.push_back(op);
   });
 
   equivalenceAnalysis(ops, aliasInfo, state);
 }
 
 /// Assert that the current bufferization decisions are consistent.

mlir/lib/Dialect/Linalg/Transforms/BufferizableOpInterfaceImpl.cpp

bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
     // Operand is read if it is used in the computation.
     auto genericOp = cast<linalg::LinalgOp>(op);
     return genericOp.payloadUsesValueFromOperand(&opOperand);
   }
 
   bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
                                const AnalysisState &state) const {
     // Operand is written to if it has an aliasing OpResult.
-    auto bufferizableOp = cast<BufferizableOpInterface>(op);
-    return !bufferizableOp.getAliasingOpResult(opOperand, state).empty();
+    auto dpsOp = cast<DestinationStyleOpInterface>(op);
+    return dpsOp.isDpsInit(&opOperand);
   }
 
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
                           const BufferizationOptions &options) const {
     return bufferizeDestinationStyleOpInterface(
         rewriter, cast<DestinationStyleOpInterface>(op), options);
   }
 };

mlir/lib/Dialect/SCF/Transforms/BufferizableOpInterfaceImpl.cpp

bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
     return true;
   }
 
   bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
                                const AnalysisState &state) const {
     return false;
   }
 
-  SmallVector<OpResult> getAliasingOpResult(Operation *op, OpOperand &opOperand,
+  AliasingOpResultList getAliasingOpResults(Operation *op, OpOperand &opOperand,
                                             const AnalysisState &state) const {
     return {};
   }
 
   bool mustBufferizeInPlace(Operation *op, OpOperand &opOperand,
                             const AnalysisState &state) const {
     // Condition operands always bufferize inplace. Otherwise, an alloc + copy
     // may be generated inside the block. We should not return/yield allocations
struct ConditionOpInterface
   }
 };
 
 /// Bufferization of scf.execute_region. Can be analyzed, but bufferization not
 /// fully implemented at the moment.
 struct ExecuteRegionOpInterface
     : public BufferizableOpInterface::ExternalModel<ExecuteRegionOpInterface,
                                                     scf::ExecuteRegionOp> {
-  SmallVector<OpOperand *>
-  getAliasingOpOperand(Operation *op, OpResult opResult,
+  AliasingOpOperandList
+  getAliasingOpOperands(Operation *op, OpResult opResult,
                         const AnalysisState &state) const {
     // ExecuteRegionOps do not have tensor OpOperands. The yielded value can be
     // any SSA value that is in scope. To allow for use-def chain traversal
     // through ExecuteRegionOps in the analysis, the corresponding yield value
     // is considered to be aliasing with the result.
     auto executeRegionOp = cast<scf::ExecuteRegionOp>(op);
     size_t resultNum = std::distance(op->getOpResults().begin(),
                                      llvm::find(op->getOpResults(), opResult));
     // TODO: Support multiple blocks.
     assert(executeRegionOp.getRegion().getBlocks().size() == 1 &&
            "expected exactly 1 block");
     auto yieldOp = dyn_cast<scf::YieldOp>(
         executeRegionOp.getRegion().front().getTerminator());
     assert(yieldOp && "expected scf.yield terminator in scf.execute_region");
-    return {&yieldOp->getOpOperand(resultNum)};
+    return {Certainty::MustBe,
+            {{&yieldOp->getOpOperand(resultNum), BufferRelation::Equivalent}}};
   }
 
   // TODO: For better bufferization results, this could return `true` only if
   // there is a memory write in the region.
   bool isMemoryWrite(Operation *op, OpResult opResult,
                      const AnalysisState &state) const {
     // Similar to scf.if, results of this op are always considered memory writes
     // in the analysis. This is a useful pattern for all ops that have tensor
for (const auto &it : llvm::enumerate(executeRegionOp->getResultTypes())) {
       }
     }
 
     // Replace old op.
     rewriter.replaceOp(executeRegionOp, newResults);
 
     return success();
   }
-
-  BufferRelation bufferRelation(Operation *op, OpResult opResult,
-                                const AnalysisState &state) const {
-    return BufferRelation::Equivalent;
-  }
 };
 
 /// Bufferization of scf.if. Replace with a new scf.if that yields memrefs.
 struct IfOpInterface
     : public BufferizableOpInterface::ExternalModel<IfOpInterface, scf::IfOp> {
-  SmallVector<OpOperand *>
-  getAliasingOpOperand(Operation *op, OpResult opResult,
+  AliasingOpOperandList
+  getAliasingOpOperands(Operation *op, OpResult opResult,
                         const AnalysisState &state) const {
     // IfOps do not have tensor OpOperands. The yielded value can be any SSA
     // value that is in scope. To allow for use-def chain traversal through
     // IfOps in the analysis, both corresponding yield values from the then/else
     // branches are considered to be aliasing with the result.
     auto ifOp = cast<scf::IfOp>(op);
     size_t resultNum = std::distance(op->getOpResults().begin(),
                                      llvm::find(op->getOpResults(), opResult));
-    return {&ifOp.thenYield()->getOpOperand(resultNum),
-            &ifOp.elseYield()->getOpOperand(resultNum)};
+    OpOperand *thenOperand = &ifOp.thenYield()->getOpOperand(resultNum);
+    OpOperand *elseOperand = &ifOp.elseYield()->getOpOperand(resultNum);
+    bool equivalentYields = state.areEquivalentBufferizedValues(
+        thenOperand->get(), elseOperand->get());
+    return {equivalentYields ? Certainty::MustBe : Certainty::MayBe,
+            {{thenOperand, BufferRelation::Equivalent},
+             {elseOperand, BufferRelation::Equivalent}}};
   }
 
   // TODO: For better bufferization results, this could return `true` only if
   // there is a memory write in one (or both) of the branches. Since this is not
   // allowed at the moment, we should never encounter scf.ifs that yield
   // unmodified tensors. Such scf.yield ops could just fold away.
   bool isMemoryWrite(Operation *op, OpResult opResult,
                      const AnalysisState &state) const {
getBufferType(Operation *op, Value value, const BufferizationOptions &options,
     // Memory space mismatch.
     if (thenBufferType.getMemorySpace() != elseBufferType.getMemorySpace())
       return op->emitError("inconsistent memory space on then/else branches");
 
     // Layout maps are different: Promote to fully dynamic layout map.
     return getMemRefTypeWithFullyDynamicLayout(
         opResult.getType().cast<TensorType>(), thenBufferType.getMemorySpace());
   }
-
-  BufferRelation bufferRelation(Operation *op, OpResult opResult,
-                                const AnalysisState &state) const {
-    // IfOp results are equivalent to their corresponding yield values if both
-    // yield values are equivalent to each other.
-    auto bufferizableOp = cast<BufferizableOpInterface>(op);
-    SmallVector<OpOperand *> yieldValues =
-        bufferizableOp.getAliasingOpOperand(opResult, state);
-    assert(yieldValues.size() == 2 && "expected 2 yield values");
-    bool equivalentYields = state.areEquivalentBufferizedValues(
-        yieldValues[0]->get(), yieldValues[1]->get());
-    return equivalentYields ? BufferRelation::Equivalent : BufferRelation::None;
-  }
 };
 
 /// Helper function for loop bufferization. Return the indices of all values
 /// that have a tensor type.
 static DenseSet<int64_t> getTensorIndices(ValueRange values) {
   DenseSet<int64_t> result;
   for (const auto &it : llvm::enumerate(values))
     if (it.value().getType().isa<TensorType>())
struct ForOpInterface
   }
 
   bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
                                const AnalysisState &state) const {
     // Tensor iter_args of scf::ForOps are always considered as a write.
     return true;
   }
 
-  SmallVector<OpResult> getAliasingOpResult(Operation *op, OpOperand &opOperand,
+  AliasingOpResultList getAliasingOpResults(Operation *op, OpOperand &opOperand,
                                             const AnalysisState &state) const {
     auto forOp = cast<scf::ForOp>(op);
-    return {forOp.getResultForOpOperand(opOperand)};
+    OpResult opResult = forOp.getResultForOpOperand(opOperand);
+    BufferRelation relation = bufferRelation(op, opResult, state);
+    return {relation == BufferRelation::Equivalent ? Certainty::MustBe
+                                                   : Certainty::MayBe,
+            {{opResult, relation}}};
   }
 
   BufferRelation bufferRelation(Operation *op, OpResult opResult,
                                 const AnalysisState &state) const {
     // ForOp results are equivalent to their corresponding init_args if the
     // corresponding iter_args and yield values are equivalent.
     auto forOp = cast<scf::ForOp>(op);
     OpOperand &forOperand = forOp.getOpOperandForResult(opResult);
struct WhileOpInterface
   }
 
   bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
                                const AnalysisState &state) const {
     // Tensor iter_args of scf::WhileOps are always considered as a write.
     return true;
   }
 
-  SmallVector<OpResult> getAliasingOpResult(Operation *op, OpOperand &opOperand,
+  AliasingOpResultList getAliasingOpResults(Operation *op, OpOperand &opOperand,
                                             const AnalysisState &state) const {
     auto whileOp = cast<scf::WhileOp>(op);
     unsigned int idx = opOperand.getOperandNumber();
 
     // The OpResults and OpOperands may not match. They may not even have the
     // same type. The number of OpResults and OpOperands can also differ.
     if (idx >= op->getNumResults() ||
         opOperand.get().getType() != op->getResult(idx).getType())
       return {};
 
     // The only aliasing OpResult may be the one at the same index.
-    return {whileOp->getResult(idx)};
+    OpResult opResult = whileOp->getResult(idx);
+    BufferRelation relation = bufferRelation(op, opResult, state);
+    return {relation == BufferRelation::Equivalent ? Certainty::MustBe
+                                                   : Certainty::MayBe,
+            {{opResult, relation}}};
   }
 
   BufferRelation bufferRelation(Operation *op, OpResult opResult,
                                 const AnalysisState &state) const {
     // WhileOp results are equivalent to their corresponding init_args if the
     // corresponding iter_args and yield values are equivalent (for both the
     // "before" and the "after" block).
     unsigned int resultNumber = opResult.getResultNumber();
bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
     return true;
   }
 
   bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
                                const AnalysisState &state) const {
     return false;
   }
 
-  SmallVector<OpResult> getAliasingOpResult(Operation *op, OpOperand &opOperand,
+  AliasingOpResultList getAliasingOpResults(Operation *op, OpOperand &opOperand,
                                             const AnalysisState &state) const {
-    if (isa<scf::IfOp>(op->getParentOp()))
-      return {op->getParentOp()->getResult(opOperand.getOperandNumber())};
+    if (auto ifOp = dyn_cast<scf::IfOp>(op->getParentOp())) {
+      Region &otherCase =
+          ifOp->getRegion((op->getParentRegion()->getRegionNumber() + 1) % 2);
+      auto otherYieldOp = cast<scf::YieldOp>(otherCase.back().back());
+      bool equivalentYields = state.areEquivalentBufferizedValues(
+          opOperand.get(),
+          otherYieldOp->getOperand(opOperand.getOperandNumber()));
+      return {equivalentYields ? Certainty::MustBe : Certainty::MayBe,
+              {{op->getParentOp()->getResult(opOperand.getOperandNumber()),
+                BufferRelation::Equivalent}}};
+    }
     if (isa<scf::ExecuteRegionOp>(op->getParentOp()))
-      return {op->getParentOp()->getResult(opOperand.getOperandNumber())};
+      return {Certainty::MustBe,
+              {{op->getParentOp()->getResult(opOperand.getOperandNumber()),
+                BufferRelation::Equivalent}}};
     return {};
   }
 
   bool mustBufferizeInPlace(Operation *op, OpOperand &opOperand,
                             const AnalysisState &state) const {
     // Yield operands always bufferize inplace. Otherwise, an alloc + copy
     // may be generated inside the block. We should not return/yield allocations
     // when possible.
struct ForeachThreadOpInterface
   }
 
   bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
                                const AnalysisState &state) const {
     // Outputs of scf::ForeachThreadOps are always considered as a write.
     return true;
   }
 
-  SmallVector<OpResult> getAliasingOpResult(Operation *op, OpOperand &opOperand,
+  AliasingOpResultList getAliasingOpResults(Operation *op, OpOperand &opOperand,
                                             const AnalysisState &state) const {
     auto foreachThreadOp = cast<ForeachThreadOp>(op);
-    return {foreachThreadOp.getTiedOpResult(&opOperand)};
-  }
-
-  BufferRelation bufferRelation(Operation *op, OpResult opResult,
-                                const AnalysisState &state) const {
-    return BufferRelation::Equivalent;
+    return {Certainty::MustBe,
+            {{foreachThreadOp.getTiedOpResult(&opOperand),
+              BufferRelation::Equivalent}}};
   }
 
   bool isWritable(Operation *op, Value value,
                   const AnalysisState &state) const {
     return true;
   }
 
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,

mlir/lib/Dialect/Shape/Transforms/BufferizableOpInterfaceImpl.cpp

 namespace mlir {
 namespace shape {
 namespace {
 
 /// Bufferization of shape.assuming.
 struct AssumingOpInterface
     : public BufferizableOpInterface::ExternalModel<AssumingOpInterface,
                                                     shape::AssumingOp> {
-  SmallVector<OpOperand *>
-  getAliasingOpOperand(Operation *op, OpResult opResult,
+  AliasingOpOperandList
+  getAliasingOpOperands(Operation *op, OpResult opResult,
                         const AnalysisState &state) const {
     // AssumingOps do not have tensor OpOperands. The yielded value can be any
     // SSA value that is in scope. To allow for use-def chain traversal through
     // AssumingOps in the analysis, the corresponding yield value is considered
     // to be aliasing with the result.
     auto assumingOp = cast<shape::AssumingOp>(op);
     size_t resultNum = std::distance(op->getOpResults().begin(),
                                      llvm::find(op->getOpResults(), opResult));
     // TODO: Support multiple blocks.
     assert(assumingOp.getDoRegion().getBlocks().size() == 1 &&
            "expected exactly 1 block");
     auto yieldOp = dyn_cast<shape::AssumingYieldOp>(
         assumingOp.getDoRegion().front().getTerminator());
     assert(yieldOp && "expected shape.assuming_yield terminator");
-    return {&yieldOp->getOpOperand(resultNum)};
+    return {Certainty::MustBe,
+            {{&yieldOp->getOpOperand(resultNum), BufferRelation::Equivalent}}};
   }
 
   // TODO: For better bufferization results, this could return `true` only if
   // there is a memory write in the region.
   bool isMemoryWrite(Operation *op, OpResult opResult,
                      const AnalysisState &state) const {
     // Similar to scf.if, results of this op are always considered memory writes
     // in the analysis. This is a useful pattern for all ops that have tensor
for (const auto &it : llvm::enumerate(assumingOp->getResultTypes())) {
       }
     }
 
     // Replace old op.
     rewriter.replaceOp(assumingOp, newResults);
 
     return success();
   }
-
-  BufferRelation bufferRelation(Operation *op, OpResult opResult,
-                                const AnalysisState &state) const {
-    return BufferRelation::Equivalent;
-  }
 };
 
 /// Bufferization of shape.assuming_yield. Bufferized as part of their enclosing
 /// ops, so this is for analysis only.
 struct AssumingYieldOpInterface
     : public BufferizableOpInterface::ExternalModel<AssumingYieldOpInterface,
                                                     shape::AssumingYieldOp> {
   bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
                               const AnalysisState &state) const {
     return true;
   }
 
   bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
                                const AnalysisState &state) const {
     return false;
   }
 
-  SmallVector<OpResult> getAliasingOpResult(Operation *op, OpOperand &opOperand,
+  AliasingOpResultList getAliasingOpResults(Operation *op, OpOperand &opOperand,
                                             const AnalysisState &state) const {
     assert(isa<shape::AssumingOp>(op->getParentOp()) &&
            "expected that parent is an AssumingOp");
-    return {op->getParentOp()->getResult(opOperand.getOperandNumber())};
+    OpResult opResult =
+        op->getParentOp()->getResult(opOperand.getOperandNumber());
+    return {Certainty::MustBe, {{opResult, BufferRelation::Equivalent}}};
   }
 
   bool mustBufferizeInPlace(Operation *op, OpOperand &opOperand,
                             const AnalysisState &state) const {
     // Yield operands always bufferize inplace. Otherwise, an alloc + copy
     // may be generated inside the block. We should not return/yield allocations
     // when possible.
     return true;

mlir/lib/Dialect/SparseTensor/Transforms/BufferizableOpInterfaceImpl.cpp

bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
     return true;
   }
 
   bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
                                const AnalysisState &state) const {
     return false;
   }
 
-  SmallVector<OpResult> getAliasingOpResult(Operation *op, OpOperand &opOperand,
+  AliasingOpResultList getAliasingOpResults(Operation *op, OpOperand &opOperand,
                                             const AnalysisState &state) const {
     return {};
   }
 
   bool isWritable(Operation *op, Value value,
                   const AnalysisState &state) const {
     return true;
   }
bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
     return true;
   }
 
   bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
                                const AnalysisState &state) const {
     return false;
   }
 
-  SmallVector<OpResult> getAliasingOpResult(Operation *op, OpOperand &opOperand,
+  AliasingOpResultList getAliasingOpResults(Operation *op, OpOperand &opOperand,
                                             const AnalysisState &state) const {
     return {};
   }
 
   bool isWritable(Operation *op, Value value,
                   const AnalysisState &state) const {
     return true;
   }
 };
 
 struct LoadOpInterface
     : public BufferizableOpInterface::ExternalModel<LoadOpInterface,
                                                     sparse_tensor::LoadOp> {
   bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
                               const AnalysisState &state) const {
     return false;
   }
 
   bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
                                const AnalysisState &state) const {
     return false;
   }
 
-  SmallVector<OpResult> getAliasingOpResult(Operation *op, OpOperand &opOperand,
+  AliasingOpResultList getAliasingOpResults(Operation *op, OpOperand &opOperand,
                                             const AnalysisState &state) const {
-    return {op->getOpResult(0)};
-  }
-
-  BufferRelation bufferRelation(Operation *op, OpResult opResult,
-                                const AnalysisState &state) const {
-    return BufferRelation::Equivalent;
+    return {Certainty::MustBe,
+            {{op->getOpResult(0), BufferRelation::Equivalent}}};
   }
 };
 
 struct NewOpInterface
     : public BufferizableOpInterface::ExternalModel<NewOpInterface,
                                                     sparse_tensor::NewOp> {
   bool isMemoryWrite(Operation *op, OpResult opResult,
                      const AnalysisState &state) const {
struct InsertOpInterface
   }
 
   bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
                                const AnalysisState &state) const {
     // InsertOp writes to memory.
     return true;
   }
 
-  SmallVector<OpResult> getAliasingOpResult(Operation *op, OpOperand &opOperand,
+  AliasingOpResultList getAliasingOpResults(Operation *op, OpOperand &opOperand,
                                             const AnalysisState &state) const {
     // InsertOp returns an alias of its operand.
     assert(op->getNumResults() == 1);
-    return op->getResults();
-  }
-
-  BufferRelation bufferRelation(Operation *oo, OpResult opResult,
-                                const AnalysisState &state) const {
-    // InsertOp returns the same object (realloc should not invalidate
-    // aliases).
-    return BufferRelation::Equivalent;
+    return {Certainty::MustBe,
+            {{op->getOpResult(0), BufferRelation::Equivalent}}};
   }
 };
 
 struct NumberOfEntriesOpInterface
     : public BufferizableOpInterface::ExternalModel<
           NumberOfEntriesOpInterface, sparse_tensor::NumberOfEntriesOp> {
   bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
                               const AnalysisState &state) const {
     return true;
   }
 
   bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
                                const AnalysisState &state) const {
     return false;
   }
 
-  SmallVector<OpResult> getAliasingOpResult(Operation *op, OpOperand &opOperand,
+  AliasingOpResultList getAliasingOpResults(Operation *op, OpOperand &opOperand,
                                             const AnalysisState &state) const {
     return {};
   }
 };
 
 struct ToIndicesOpInterface
     : public BufferizableOpInterface::ExternalModel<
           ToIndicesOpInterface, sparse_tensor::ToIndicesOp> {
   bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
                               const AnalysisState &state) const {
     return true;
   }
 
   bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
                                const AnalysisState &state) const {
     // Potential writes into memory through the result of sparse_tensor.indices
     // are not considered.
     return false;
   }
 
-  SmallVector<OpResult> getAliasingOpResult(Operation *op, OpOperand &opOperand,
+  AliasingOpResultList getAliasingOpResults(Operation *op, OpOperand &opOperand,
                                             const AnalysisState &state) const {
     return {};
   }
 };
 
 struct ToPointersOpInterface
     : public BufferizableOpInterface::ExternalModel<
           ToPointersOpInterface, sparse_tensor::ToPointersOp> {
   bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
                               const AnalysisState &state) const {
     return true;
   }
 
   bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
                                const AnalysisState &state) const {
     // Potential writes into memory through the result of sparse_tensor.pointers
     // are not considered.
     return false;
   }
 
-  SmallVector<OpResult> getAliasingOpResult(Operation *op, OpOperand &opOperand,
+  AliasingOpResultList getAliasingOpResults(Operation *op, OpOperand &opOperand,
                                             const AnalysisState &state) const {
     return {};
   }
 };
 
 struct ToValuesOpInterface
     : public BufferizableOpInterface::ExternalModel<ToValuesOpInterface,
                                                     sparse_tensor::ToValuesOp> {
   bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
                               const AnalysisState &state) const {
     return true;
   }
 
   bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
                                const AnalysisState &state) const {
     // Potential writes into memory through the result of sparse_tensor.values
     // are not considered.
     return false;
   }
 
-  SmallVector<OpResult> getAliasingOpResult(Operation *op, OpOperand &opOperand,
+  AliasingOpResultList getAliasingOpResults(Operation *op, OpOperand &opOperand,
                                             const AnalysisState &state) const {
     return {};
   }
 };
 
 } // namespace
 } // namespace sparse_tensor
 } // namespace mlir

mlir/lib/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.cpp

bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
     return false;
   }
 
   bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
                                const AnalysisState &state) const {
     return false;
   }
 
-  SmallVector<OpResult> getAliasingOpResult(Operation *op, OpOperand &opOperand,
+  AliasingOpResultList getAliasingOpResults(Operation *op, OpOperand &opOperand,
                                             const AnalysisState &state) const {
-    return {op->getResult(0)};
-  }
-
-  BufferRelation bufferRelation(Operation *op, OpResult opResult,
-                                const AnalysisState &state) const {
-    return BufferRelation::Equivalent;
+    return {Certainty::MustBe,
+            {{op->getResult(0), BufferRelation::Equivalent}}};
   }
 
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
                           const BufferizationOptions &options) const {
     auto castOp = cast<tensor::CastOp>(op);
 
     // The result buffer still has the old (pre-cast) type.
     FailureOr<Value> resultBuffer =
bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
     return false;
   }
 
   bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
                                const AnalysisState &state) const {
     return false;
   }
 
-  SmallVector<OpResult> getAliasingOpResult(Operation *op, OpOperand &opOperand,
-                                            const AnalysisState &state) const {
-    if (&opOperand == &op->getOpOperand(0) /*src*/)
-      return {op->getOpResult(0)};
-    return {};
-  }
-
-  BufferRelation bufferRelation(Operation *op, OpResult opResult,
+  AliasingOpResultList getAliasingOpResults(Operation *op, OpOperand &opOperand,
                                             const AnalysisState &state) const {
-    return BufferRelation::Equivalent;
+    // CollapseShapeOp may allocate at runtime.
+    // TODO: Predict future buffer type and compute more accurate Certainty.
+    return {Certainty::MayBe,
+            {{op->getOpResult(0), BufferRelation::Equivalent}}};
   }
 
   FailureOr<BaseMemRefType>
   getBufferType(Operation *op, Value value, const BufferizationOptions &options,
                 const DenseMap<Value, BaseMemRefType> &fixedTypes) const {
     auto collapseShapeOp = cast<tensor::CollapseShapeOp>(op);
     auto maybeSrcBufferType = bufferization::getBufferType(
         collapseShapeOp.getSrc(), options, fixedTypes);
bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
     return false;
   }
 
   bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
                                const AnalysisState &state) const {
     return false;
   }
 
-  SmallVector<OpResult> getAliasingOpResult(Operation *op, OpOperand &opOperand,
+  AliasingOpResultList getAliasingOpResults(Operation *op, OpOperand &opOperand,
                                             const AnalysisState &state) const {
     return {};
   }
 
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
                           const BufferizationOptions &options) const {
     auto dimOp = cast<tensor::DimOp>(op);
     FailureOr<Value> v = getBuffer(rewriter, dimOp.getSource(), options);
bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
     return false;
   }
 
   bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
                                const AnalysisState &state) const {
     return false;
   }
 
-  SmallVector<OpResult> getAliasingOpResult(Operation *op, OpOperand &opOperand,
-                                            const AnalysisState &state) const {
-    if (&opOperand == &op->getOpOperand(0) /*src*/)
-      return {op->getOpResult(0)};
-    return {};
-  }
-
-  BufferRelation bufferRelation(Operation *op, OpResult opResult,
+  AliasingOpResultList getAliasingOpResults(Operation *op, OpOperand &opOperand,
                                             const AnalysisState &state) const {
-    return BufferRelation::Equivalent;
+    return {Certainty::MustBe,
+            {{op->getOpResult(0), BufferRelation::Equivalent}}};
   }
 
   FailureOr<BaseMemRefType>
   getBufferType(Operation *op, Value value, const BufferizationOptions &options,
                 const DenseMap<Value, BaseMemRefType> &fixedTypes) const {
     auto expandShapeOp = cast<tensor::ExpandShapeOp>(op);
     auto maybeSrcBufferType = bufferization::getBufferType(
         expandShapeOp.getSrc(), options, fixedTypes);
bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
     return false;
   }
 
   bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
                                const AnalysisState &state) const {
     return false;
   }
 
-  SmallVector<OpResult> getAliasingOpResult(Operation *op, OpOperand &opOperand,
-                                            const AnalysisState &state) const {
-    if (&opOperand == &op->getOpOperand(0) /*source*/)
-      return {op->getOpResult(0)};
-    return {};
-  }
-
-  BufferRelation bufferRelation(Operation *op, OpResult opResult,
+  AliasingOpResultList getAliasingOpResults(Operation *op, OpOperand &opOperand,
                                             const AnalysisState &state) const {
-    return BufferRelation::None;
+    return {Certainty::MustBe, {{op->getOpResult(0), BufferRelation::None}}};
   }
 
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
                           const BufferizationOptions &options) const {
     auto extractSliceOp = cast<tensor::ExtractSliceOp>(op);
     SmallVector<OpFoldResult> mixedOffsets = extractSliceOp.getMixedOffsets();
     SmallVector<OpFoldResult> mixedSizes = extractSliceOp.getMixedSizes();
     SmallVector<OpFoldResult> mixedStrides = extractSliceOp.getMixedStrides();
bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
     return true;
   }
 
   bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
                                const AnalysisState &state) const {
     return false;
   }
 
-  SmallVector<OpResult> getAliasingOpResult(Operation *op, OpOperand &opOperand,
+  AliasingOpResultList getAliasingOpResults(Operation *op, OpOperand &opOperand,
                                             const AnalysisState &state) const {
     return {};
   }
 
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
                           const BufferizationOptions &options) const {
     auto extractOp = cast<tensor::ExtractOp>(op);
     FailureOr<Value> srcMemref =
bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
     return true;
   }
 
   bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
                                const AnalysisState &state) const {
     return false;
   }
 
-  SmallVector<OpResult> getAliasingOpResult(Operation *op, OpOperand &opOperand,
+  AliasingOpResultList getAliasingOpResults(Operation *op, OpOperand &opOperand,
                                             const AnalysisState &state) const {
     return {};
   }
 
   FailureOr<BaseMemRefType>
   getBufferType(Operation *op, Value value, const BufferizationOptions &options,
                 const DenseMap<Value, BaseMemRefType> &fixedTypes) const {
     // Infer memory space from the source tensor.
bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
     return false;
   }
 
   bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
                                const AnalysisState &state) const {
     return false;
   }
 
-  SmallVector<OpResult> getAliasingOpResult(Operation *op, OpOperand &opOperand,
+  AliasingOpResultList getAliasingOpResults(Operation *op, OpOperand &opOperand,
                                             const AnalysisState &state) const {
     return {};
   }
 
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
                           const BufferizationOptions &options) const {
     auto rankOp = cast<tensor::RankOp>(op);
     FailureOr<Value> v = getBuffer(rewriter, rankOp.getTensor(), options);
bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
     return false;
   }
 
   bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
                                const AnalysisState &state) const {
     return false;
   }
 
-  SmallVector<OpResult> getAliasingOpResult(Operation *op, OpOperand &opOperand,
-                                            const AnalysisState &state) const {
-    return {op->getOpResult(0)};
-  }
-
-  BufferRelation bufferRelation(Operation *op, OpResult opResult,
+  AliasingOpResultList getAliasingOpResults(Operation *op, OpOperand &opOperand,
                                             const AnalysisState &state) const {
-    return BufferRelation::Equivalent;
+    return {Certainty::MustBe,
+            {{op->getOpResult(0), BufferRelation::Equivalent}}};
   }
 
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
                           const BufferizationOptions &options) const {
     auto reshapeOp = cast<tensor::ReshapeOp>(op);
     FailureOr<Value> srcBuffer =
         getBuffer(rewriter, reshapeOp.getSource(), options);
     FailureOr<Value> shapeBuffer =
         getBuffer(rewriter, reshapeOp.getShape(), options);
     if (failed(srcBuffer) || failed(shapeBuffer))
       return failure();
     auto resultMemRefType = getMemRefType(
         reshapeOp.getResult(), options, /*layout=*/{},
         srcBuffer->getType().cast<BaseMemRefType>().getMemorySpace());
     replaceOpWithNewBufferizedOp<memref::ReshapeOp>(
         rewriter, op, resultMemRefType, *srcBuffer, *shapeBuffer);
     return success();
   }
 };
 
 /// Analysis of ParallelInsertSliceOp.
 struct ParallelInsertSliceOpInterface
     : public BufferizableOpInterface::ExternalModel<
           ParallelInsertSliceOpInterface, ParallelInsertSliceOp> {
-  SmallVector<OpResult> getAliasingOpResult(Operation *op, OpOperand &opOperand,
+  AliasingOpResultList getAliasingOpResults(Operation *op, OpOperand &opOperand,
                                             const AnalysisState &state) const {
-      return {};
   }
 
   bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
                               const AnalysisState &state) const {
     return true;
   }
 
   bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,

mlir/lib/Dialect/Vector/Transforms/BufferizableOpInterfaceImpl.cpp

struct TransferReadOpInterface
 
   bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
                                const AnalysisState &state) const {
     assert(opOperand.get().getType().isa<RankedTensorType>() &&
            "only tensor types expected");
     return false;
   }
 
-  SmallVector<OpResult> getAliasingOpResult(Operation *op, OpOperand &opOperand,
+  AliasingOpResultList getAliasingOpResults(Operation *op, OpOperand &opOperand,
                                             const AnalysisState &state) const {
     return {};
   }
 
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
                           const BufferizationOptions &options) const {
     auto readOp = cast<vector::TransferReadOp>(op);
     assert(readOp.getShapedType().isa<TensorType>() &&
struct GatherOpInterface
 
   bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
                                const AnalysisState &state) const {
     assert(opOperand.get().getType().isa<RankedTensorType>() &&
            "only tensor types expected");
     return false;
   }
 
-  SmallVector<OpResult> getAliasingOpResult(Operation *op, OpOperand &opOperand,
+  AliasingOpResultList getAliasingOpResults(Operation *op, OpOperand &opOperand,
                                             const AnalysisState &state) const {
     return {};
   }
 
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
                           const BufferizationOptions &options) const {
     auto gatherOp = cast<vector::GatherOp>(op);
     assert(gatherOp.getBaseType().isa<TensorType>() &&