[mlir][bufferization] Generalize getAliasingOpResults to getAliasingValues

This revision is needed to support bufferization of `cf.br`/`cf.cond_br`. It will also be useful for better analysis of loop ops.

This revision generalizes `getAliasingOpResults` to `getAliasingValues`. An OpOperand can now not only alias with OpResults but also with BlockArguments. In the case of `cf.br` (will be added in a later revision): a `cf.br` operand will alias with the corresponding argument of the destination block.

If an op does not implement the `BufferizableOpInterface`, the analysis in conservative. It previously assumed that an OpOperand may alias with each OpResult. It now assumes that an OpOperand may alias with each OpResult and each BlockArgument of the entry block.

Differential Revision: https://reviews.llvm.org/D157957

Author: matthias-springer

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

@@ -49,14 +49,13 @@ struct AliasingOpOperand {
   bool isDefinite;
 };
 
-/// A maybe aliasing OpResult. If `isDefinite` is `true`, the OpResult is
-/// guaranteed to alias at runtime.
-struct AliasingOpResult {
-  AliasingOpResult(OpResult opResult, BufferRelation relation,
-                   bool isDefinite = true)
-      : opResult(opResult), relation(relation), isDefinite(isDefinite) {}
+/// A maybe aliasing Value. If `isDefinite` is `true`, the Value is guaranteed
+/// to alias at runtime.
+struct AliasingValue {
+  AliasingValue(Value value, BufferRelation relation, bool isDefinite = true)
+      : value(value), relation(relation), isDefinite(isDefinite) {}
 
-  OpResult opResult;
+  Value value;
   BufferRelation relation;
   bool isDefinite;
 };
@@ -90,12 +89,12 @@ template <typename T> class AliasList {
 };
 
 /// A list of possible aliasing OpOperands. This list models the runtime
-/// aliasing relationship for an OpResult.
+/// aliasing relationship for a Value.
 using AliasingOpOperandList = AliasList<AliasingOpOperand>;
 
-/// A list of possible aliasing OpResults. This list models the runtime
-/// aliasing relationship for an OpOperand.
-using AliasingOpResultList = AliasList<AliasingOpResult>;
+/// A list of possible aliasing Values. This list models the runtime aliasing
+/// relationship for an OpOperand.
+using AliasingValueList = AliasList<AliasingValue>;
 
 class OpFilter {
 public:
@@ -418,15 +417,14 @@ struct TraversalConfig {
 /// tensor values.
 class AnalysisState {
 public:
-  /// Determine which OpOperand* will alias with `result` if the op is
+  /// Determine which OpOperand* will alias with `value` if the op is
   /// bufferized in place. Return all tensor OpOperand* if the op is not
   /// bufferizable.
-  AliasingOpOperandList getAliasingOpOperands(OpResult result) const;
+  AliasingOpOperandList getAliasingOpOperands(Value value) 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.
-  AliasingOpResultList getAliasingOpResults(OpOperand &opOperand) const;
+  /// Determine which Value will alias with `opOperand` if the op is bufferized
+  /// in place. Return all tensor Values if the op is not bufferizable.
+  AliasingValueList getAliasingValues(OpOperand &opOperand) const;
 
   /// Return true if `opOperand` bufferizes to a memory read. Return `true` if
   /// the op is not bufferizable.
@@ -685,7 +683,7 @@ namespace detail {
 /// This is the default implementation of
 /// BufferizableOpInterface::getAliasingOpOperands. Should not be called from
 /// other places.
-AliasingOpOperandList defaultGetAliasingOpOperands(OpResult opResult,
+AliasingOpOperandList defaultGetAliasingOpOperands(Value value,
                                                    const AnalysisState &state);
 
 /// This is the default implementation of
@@ -709,11 +707,11 @@ bool defaultIsRepetitiveRegion(BufferizableOpInterface bufferizableOp,
 
 /// This is the default implementation of getAliasingOpOperands in case the
 /// defining op does not implement the BufferizableOpInterface.
-AliasingOpOperandList unknownGetAliasingOpOperands(OpResult opResult);
+AliasingOpOperandList unknownGetAliasingOpOperands(Value value);
 
-/// This is the default implementation of getAliasingOpResults in case the
-/// owner op does not implement the BufferizableOpInterface.
-AliasingOpResultList unknownGetAliasingOpResults(OpOperand &opOperand);
+/// This is the default implementation of getAliasingValues in case the owner
+/// op does not implement the BufferizableOpInterface.
+AliasingValueList unknownGetAliasingValues(OpOperand &opOperand);
 } // namespace detail
 
 } // namespace bufferization

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

@@ -18,23 +18,23 @@ def BufferizableOpInterface : OpInterface<"BufferizableOpInterface"> {
 
     Note: All "bufferizesTo*" and "getAliasing*" interface methods must be
     implemented conservatively. If it is not statically known whether an
-    OpOperand/OpResult bufferizes in a certain way (e.g., to a memory write),
+    OpOperand/Value bufferizes in a certain way (e.g., to a memory write),
     the worst case must be assumed (e.g., that it does). Similarly,
     "getAliasing*" interface methods may always return additional OpOperands or
-    OpResults, but must not miss an OpOperand or OpResult that could potentially
+    Values, but must not miss an OpOperand or Value that could potentially
     alias at runtime.
   }];
   let cppNamespace = "::mlir::bufferization";
   let methods = [
       InterfaceMethod<
         /*desc=*/[{
-          Return `true` if the given OpResult may bufferize to a new buffer
-          allocation. If it is statically unknown if the given OpResult
+          Return `true` if the given Value may bufferize to a new buffer
+          allocation. If it is statically unknown that the given Value
           bufferizes to a buffer allocation, `true` should be returned.
         }],
         /*retType=*/"bool",
         /*methodName=*/"bufferizesToAllocation",
-        /*args=*/(ins "::mlir::OpResult":$opResult),
+        /*args=*/(ins "::mlir::Value":$value),
         /*methodBody=*/"",
         /*defaultImplementation=*/"return false;"
       >,
@@ -68,10 +68,9 @@ def BufferizableOpInterface : OpInterface<"BufferizableOpInterface"> {
           tensor type.
 
           This method will never be called on OpOperands that do not have an
-          aliasing OpResult. Intuitively, it does not make sense for an
-          OpOperand to bufferize to a memory write without returning an aliasing
-          tensor, because the write would have no visible effect outside of the
-          op.
+          aliasing Value. Intuitively, it does not make sense for an OpOperand
+          to bufferize to a memory write without returning an aliasing tensor,
+          because the write would have no visible effect outside of the op.
 
           Note: It is always safe to consider an OpOperand as a memory write,
           even if it does actually not write; however, this can introduce
@@ -87,7 +86,7 @@ def BufferizableOpInterface : OpInterface<"BufferizableOpInterface"> {
         /*defaultImplementation=*/[{
           // Does not have to be implemented for ops without tensor OpOperands.
           // Does not have to be implemented for OpOperands that do not have an
-          // aliasing OpResult.
+          // aliasing Value.
           llvm_unreachable("bufferizesToMemoryWrite not implemented");
          }]
       >,
@@ -221,21 +220,21 @@ def BufferizableOpInterface : OpInterface<"BufferizableOpInterface"> {
       >,
       InterfaceMethod<
         /*desc=*/[{
-          Return the OpResults that may alias with a given OpOperand when
+          Return the Values 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.
 
-          This method can return multiple OpResults, indicating that a given
+          This method can return multiple Values, indicating that a given
           OpOperand may at runtime alias with any (or multiple) of the returned
-          OpResults.
+          Values.
 
           Each alias is specified with a degree of certainty:
 
           * MAYBE (`isDefinite = false`): At runtime, buffer(opOperand) may
-            alias with the specified OpResult.
+            alias with the specified Value.
           * DEFINITE (`isDefinite = true`, default): At runtime,
             buffer(opOperand) is guaranteed to alias the buffer of the specified
-            OpResult. This is a stronger property than MAYBE and allows for more
+            Value. This is a stronger property than MAYBE and allows for more
             precise analyses. DEFINITE properties should be used when possible.
 
           Furthermore, each alias is specified with a buffer relation:
@@ -245,69 +244,69 @@ def BufferizableOpInterface : OpInterface<"BufferizableOpInterface"> {
           * `BufferRelation::Unknown`: There is no further information apart
             from the fact that both buffers alias.
 
-          False positives are allowed in the list of OpResults, but they can
+          False positives are allowed in the list of Values, but they can
           adversely affect the accuracy of the anlysis. On the contrary,
           omitting potential aliases is incorrect.
 
           One possible (conservative) implementation of this interface method,
-          that is always safe, is to return all tensor OpResults with
+          that is always safe, is to return all tensor Values with
           BufferRelation::Unknown and MAYBE.
 
           Examples:
 
           ```
-          // aliasingOpResults(%t) = DEFINITE {Equivalent %r}
+          // aliasingValues(%t) = DEFINITE {Equivalent %r}
           %r = tensor.insert_slice %f into %t : tensor<10xf32>
 
-          // aliasingOpResults(%t) = DEFINITE {Unknown %r}
+          // aliasingValues(%t) = DEFINITE {Unknown %r}
           // Note: "Buffer is subset of buffer" relationship are not yet
           // supported, so "Unknown" is the best we can do for now.
           %r = tensor.extract_slice %t[0]][5][1]
               : tensor<10xf32> to tensor<5xf32>
 
-          // aliasingOpResults(%t1) = MAYBE {Equivalent %r}
-          // aliasingOpResults(%t2) = MAYBE {Equivalent %r}
+          // aliasingValues(%t1) = MAYBE {Equivalent %r}
+          // aliasingValues(%t2) = MAYBE {Equivalent %r}
           %r = arith.select %c, %t1, %t2 : tensor<10xf32>
 
           // A hypothetical op that bufferizes to rolling a dice and based on
           // the result to either return buffer(%t) or a newly allocated copy
           // thereof.
-          // aliasingOpResults(%t) = MAYBE {Equivalent %r}
+          // aliasingValues(%t) = MAYBE {Equivalent %r}
           %r = "dummy.alias_or_copy(%t) : (tensor<10xf32>) -> (tensor<10xf32>)"
           ```
         }],
-        /*retType=*/"::mlir::bufferization::AliasingOpResultList",
-        /*methodName=*/"getAliasingOpResults",
+        /*retType=*/"::mlir::bufferization::AliasingValueList",
+        /*methodName=*/"getAliasingValues",
         /*args=*/(ins "::mlir::OpOperand &":$opOperand,
                       "const ::mlir::bufferization::AnalysisState &":$state),
         /*methodBody=*/"",
         /*defaultImplementation=*/[{
           // Does not have to be implemented for ops without tensor OpOperands.
           assert(::llvm::isa<::mlir::TensorType>(opOperand.get().getType()) &&
                  "expected OpOperand with tensor type");
-          llvm_unreachable("getAliasingOpResults not implemented");
+          llvm_unreachable("getAliasingValues 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.
+          Return the OpOperands that alias with a given Value when bufferized
+          in-place. This method will never be called on Values that do not
+          have a tensor type.
 
-          By default, this method is the inverse of `getAliasingOpResults`. Ops
+          By default, this method is the inverse of `getAliasingValues`. Ops
           with a region that yield values may want to override this method to
           return the OpOperands that are yielded by the terminator.
 
           This method can return multiple OpOperands, indicating that a given
-          OpResult may at runtime alias with any (or multiple) of the returned
+          Value may at runtime alias with any (or multiple) of the returned
           OpOperands.
 
           This property is specified with a degree of certainty:
 
-          * MAYBE (`isDefinite = false`): At runtime, buffer(opResult) may
-            alias with the specified OpOperand.
+          * MAYBE (`isDefinite = false`): At runtime, buffer(value) may alias
+            with the specified OpOperand.
           * DEFINITE (`isDefinite = true`, default): At runtime,
-            buffer(opResult) is guaranteed to alias the buffer of the specified
+            buffer(value) is guaranteed to alias the buffer of the specified
             OpOperand. This is a stronger property than MAYBE and allows for
             more precise analyses. DEFINITE properties should be used when
             possible.
@@ -350,14 +349,14 @@ def BufferizableOpInterface : OpInterface<"BufferizableOpInterface"> {
         }],
         /*retType=*/"::mlir::bufferization::AliasingOpOperandList",
         /*methodName=*/"getAliasingOpOperands",
-        /*args=*/(ins "::mlir::OpResult":$opResult,
+        /*args=*/(ins "::mlir::Value":$value,
                       "const ::mlir::bufferization::AnalysisState &":$state),
         /*methodBody=*/"",
         /*defaultImplementation=*/[{
-          assert(::llvm::isa<::mlir::TensorType>(opResult.getType()) &&
-                 "expected OpResult with tensor type");
+          assert(isa<::mlir::TensorType>(value.getType()) &&
+                 "expected tensor type");
           return ::mlir::bufferization::detail::defaultGetAliasingOpOperands(
-              opResult, state);
+              value, state);
         }]
       >,
       InterfaceMethod<
@@ -568,7 +567,7 @@ def BufferizableOpInterface : OpInterface<"BufferizableOpInterface"> {
           ::llvm::cast<::mlir::bufferization::BufferizableOpInterface>(getOperation());
       return !bufferizableOp.bufferizesToMemoryRead(opOperand, state)
           && !bufferizableOp.bufferizesToMemoryWrite(opOperand, state)
-          && bufferizableOp.getAliasingOpResults(opOperand, state)
+          && bufferizableOp.getAliasingValues(opOperand, state)
               .getNumAliases() != 0;
     }
   }];

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

@@ -95,15 +95,15 @@ def Bufferization_AllocTensorOp : Bufferization_Op<"alloc_tensor",
     bool resultBufferizesToMemoryWrite(OpResult opResult,
                                        const AnalysisState &state);
 
-    bool bufferizesToAllocation(OpResult opResult) { return true; }
+    bool bufferizesToAllocation(Value value) { return true; }
 
     bool bufferizesToMemoryRead(OpOperand &opOperand,
                                 const AnalysisState &state);
 
     bool bufferizesToMemoryWrite(OpOperand &opOperand,
                                  const AnalysisState &state);
 
-    AliasingOpResultList getAliasingOpResults(
+    AliasingValueList getAliasingValues(
         OpOperand &opOperand, const AnalysisState &state);
 
     FailureOr<BaseMemRefType> getBufferType(
@@ -237,7 +237,7 @@ def Bufferization_CopyTensorOp : Bufferization_Op<"copy_tensor",
     bool bufferizesToMemoryWrite(OpOperand &opOperand,
                                  const AnalysisState &state);
 
-    AliasingOpResultList getAliasingOpResults(
+    AliasingValueList getAliasingValues(
         OpOperand &opOperand, const AnalysisState &state);
 
     RankedTensorType getType() {
@@ -295,7 +295,7 @@ def Bufferization_DeallocTensorOp : Bufferization_Op<"dealloc_tensor",
       return false;
     }
 
-    AliasingOpResultList getAliasingOpResults(
+    AliasingValueList getAliasingValues(
         OpOperand &opOperand, const AnalysisState &state) const {
       return {};
     }
@@ -459,7 +459,7 @@ def Bufferization_ToMemrefOp : Bufferization_Op<"to_memref", [
       return !getReadOnly();
     }
 
-    AliasingOpResultList getAliasingOpResults(
+    AliasingValueList getAliasingValues(
         OpOperand &opOperand, const AnalysisState &state) const {
       return {};
     }

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

@@ -36,8 +36,8 @@ struct DstBufferizableOpInterfaceExternalModel
     return dstOp.isDpsInit(&opOperand);
   }
 
-  AliasingOpResultList getAliasingOpResults(Operation *op, OpOperand &opOperand,
-                                            const AnalysisState &state) const {
+  AliasingValueList getAliasingValues(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))

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

@@ -76,8 +76,8 @@ struct IndexCastOpInterface
     return false;
   }
 
-  AliasingOpResultList getAliasingOpResults(Operation *op, OpOperand &opOperand,
-                                            const AnalysisState &state) const {
+  AliasingValueList getAliasingValues(Operation *op, OpOperand &opOperand,
+                                      const AnalysisState &state) const {
     return {{op->getResult(0), BufferRelation::Equivalent}};
   }
 
@@ -123,8 +123,8 @@ struct SelectOpInterface
     return false;
   }
 
-  AliasingOpResultList getAliasingOpResults(Operation *op, OpOperand &opOperand,
-                                            const AnalysisState &state) const {
+  AliasingValueList getAliasingValues(Operation *op, OpOperand &opOperand,
+                                      const AnalysisState &state) const {
     return {{op->getOpResult(0) /*result*/, BufferRelation::Equivalent,
              /*isDefinite=*/false}};
   }