[mlir][bufferization][NFC] Make `escape` a dialect attribute

[springerm]

Summary

All bufferizable ops that bufferize to an allocation receive a bufferization.escape attribute during TensorCopyInsertion.

[springerm]

update

[aartbik]

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

31

did you really intent, "may bufferize to", or will "bufferizes to" be better?

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

43

perhaps add one more sentence on clarifying "escape", for example

which means the buffer needs to be allocated outside this method

or

which means the allocation "escape" the scope of the method to the caller

(or something like that

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

69

Can't we use this condition to guard the whole analysis?

[springerm]

update

[springerm]

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

31

Yes, this is on purpose. Sometimes we don't know what's going to happen at runtime. Then we should return true here. Extended the comment.

[springerm]

address comments

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

 #define BUFFERIZABLE_OP_INTERFACE
 
 include "mlir/IR/OpBase.td"
 
 def BufferizableOpInterface : OpInterface<"BufferizableOpInterface"> {
   let description = [{
     An op interface for One-Shot Bufferize. Ops that implement this interface
     interface can be analyzed and bufferized using One-Shot Bufferize.
+
+    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),
+    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
+    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
+          bufferizes to a buffer allocation, `true` should be returned.
+        }],
+        /*retType=*/"bool",
+        /*methodName=*/"bufferizesToAllocation",
+        /*args=*/(ins "OpResult":$opResult),
+        /*methodBody=*/"",
+        /*defaultImplementation=*/"return false;"
+      >,
+      InterfaceMethod<
+        /*desc=*/[{
           Return `true` if the given OpOperand bufferizes to a memory read. This
           method will never be called on OpOperands that do not have a tensor
           type.
 
           Note: It is always safe to consider an OpOperand as a memory read,
           even if it does actually not read; however, this can introduce
           unnecessary out-of-place bufferization decisions. One-Shot Analysis
           considers OpOperands of unknown ops (that do not implement this

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

let extraClassDeclaration = [{
     /// arguments during One-Shot Module Bufferize.
     constexpr const static ::llvm::StringLiteral
         kWritableAttrName = "bufferization.writable";
 
     /// Attribute name used to mark the bufferization layout for region
     /// arguments during One-Shot Module Bufferize.
     constexpr const static ::llvm::StringLiteral
         kBufferLayoutAttrName = "bufferization.buffer_layout";
+
+    /// Attribute name used to mark escaping behavior of buffer allocations.
+    /// Escaping allocations cannot be deallocated in the same block and must
+    /// be treated specially: They are currently deallocated with the
+    /// BufferDeallocation pass.
+    ///
+    /// Note: Only ops with at least one OpResult that bufferizes to a buffer
+    /// allocation (as per BufferizableOpInterface) may have this attribute.
+    constexpr const static ::llvm::StringLiteral
+        kEscapeAttrName = "bufferization.escape";
   }];
   let hasOperationAttrVerify = 1;
   let emitAccessorPrefix = kEmitAccessorPrefix_Prefixed;
 }
 
 #endif // BUFFERIZATION_BASE

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

let description = [{
     given shape (dynamic or static). It always bufferizes to a new buffer
     allocation of the given shape. The optional `copy` operand specifies the
     contents of the tensors. If no `copy` operand is specified, reading from the
     result of an `alloc_tensor` op yields an undefined value.
 
     If `copy` is specified, no dynamic sizes should be passed, since they are
     the same as the dynamic sizes of the `copy` operand.
 
-    The optional `escape` attribute indicates whether the buffer escapes the
-    parent block or not. In the latter case, the buffer is deallocated at the
-    of the block (during bufferization). In the former case, the buffer is not
-    deallocated and must be deallocated through some other mechanism.
-
     `alloc_tensor` is a helper op for bufferization. The operation is provided
     as an anchor that marks the beginning of a new tensor SSA use-def chain. It
     can be used to control in-place bufferization decisions during One-Shot
     Bufferize: The bufferized result of a `bufferization.alloc_tensor` does not
     alias with any other buffer, so it can be used to resolve read-after-write
     conflicts that would have been introduced by the in-place bufferization of
     another op.
 
     Both dense and sparse tensor types are supported. The result of a
     `bufferization.alloc_tensor` is a tensor value that can be used like any
     other tensor value. In practice, it is often used as the "out" operand of
     another op. E.g.:
 
     ```mlir
     %c = bufferization.alloc_tensor [%d1, %d2] : tensor<?x?xf32, #SparseMatrix>
     %0 = linalg.matmul
       ins(%a, %b: tensor<?x?xf32, #SparseMatrix>, tensor<?x?xf32, #SparseMatrix>)
       outs(%c: tensor<?x?xf32, #SparseMatrix>) -> tensor<?x?xf32, #SparseMatrix>
     ```
   }];
 
   let arguments = (ins Variadic<Index>:$dynamic_sizes,
-                       Optional<AnyTensor>:$copy,
-                       OptionalAttr<BoolAttr>:$escape);
+                       Optional<AnyTensor>:$copy);
 
   let results = (outs AnyTensor:$result);
 
   let extraClassDeclaration = [{
     LogicalResult bufferize(RewriterBase &rewriter,
                             const BufferizationOptions &options);
 
     bool isMemoryWrite(OpResult opResult, const AnalysisState &state);
 
+    bool bufferizesToAllocation(OpResult opResult) { return true; }
+
     bool bufferizesToMemoryRead(OpOperand &opOperand,
                                 const AnalysisState &state);
 
     bool bufferizesToMemoryWrite(OpOperand &opOperand,
                                  const AnalysisState &state);
 
     SmallVector<OpResult> getAliasingOpResult(
         OpOperand &opOperand, const AnalysisState &state);
let extraClassDeclaration = [{
     // and return the shape.
     int64_t getStaticSize(unsigned idx) {
       assert(!isDynamicDim(idx) && "expected static size");
       return getType().getShape()[idx];
     }
   }];
 
   let builders = [
-    // Build an op without `copy` operand and `escape` attribute.
+    // Build an op without `copy` operand.
     OpBuilder<(ins "RankedTensorType":$type, "ValueRange":$dynamicSizes)>,
-
-    // Build an op without `escape` attribute.
-    OpBuilder<(ins "RankedTensorType":$type, "ValueRange":$dynamicSizes,
-                   "Value":$copy)>,
-
-    // Build an op with `copy` and `escape` attribute.
-    OpBuilder<(ins "RankedTensorType":$type, "ValueRange":$dynamicSizes,
-                   "Value":$copy, "bool":$escape)>,
   ];
 
   let hasCanonicalizer = 1;
   let hasCustomAssemblyFormat = 1;
   let hasVerifier = 1;
 }
 
 //===----------------------------------------------------------------------===//

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

if (shapedValue.getType().isa<RankedTensorType>() &&
       }
     }
 
     // If the shape could not be reified, create DimOps.
     if (!reifiedShapes)
       populateDynamicDimSizes(b, loc, tensor, dynamicSizes);
   }
 
-  return b.create<AllocTensorOp>(loc, tensorType, dynamicSizes,
-                                 copy ? tensor : Value(), escape);
+  auto allocTensorOp = b.create<AllocTensorOp>(loc, tensorType, dynamicSizes,
+                                               copy ? tensor : Value());
+  allocTensorOp->setAttr(BufferizationDialect::kEscapeAttrName,
+                         b.getBoolArrayAttr({escape}));
+  return allocTensorOp;
 }
 
 LogicalResult BufferizableOpInterface::resolveTensorOpOperandConflicts(
     RewriterBase &rewriter, const AnalysisState &state) {
   OpBuilder::InsertionGuard g(rewriter);
   Operation *op = getOperation();
   SmallVector<OpOperand *> outOfPlaceOpOperands;
   DenseSet<OpOperand *> copiedOpOperands;

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

 //===----------------------------------------------------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
+#include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h"
 #include "mlir/Dialect/Bufferization/IR/Bufferization.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/IR/FunctionInterfaces.h"
 #include "mlir/Transforms/InliningUtils.h"
 
 using namespace mlir;
 using namespace mlir::bufferization;
 
 #include "mlir/Dialect/Bufferization/IR/BufferizationOpsDialect.cpp.inc"
 
 /// Attribute name used to mark function arguments who's buffers can be written
 /// to during One-Shot Module Bufferize.
 constexpr const ::llvm::StringLiteral BufferizationDialect::kWritableAttrName;
 
 /// Attribute name used to mark the bufferization layout for region arguments
 /// during One-Shot Module Bufferize.
 constexpr const ::llvm::StringLiteral
     BufferizationDialect::kBufferLayoutAttrName;
 
+/// Attribute name used to mark escaping behavior of buffer allocations.
+constexpr const ::llvm::StringLiteral BufferizationDialect::kEscapeAttrName;
+
 //===----------------------------------------------------------------------===//
 // Bufferization Dialect Interfaces
 //===----------------------------------------------------------------------===//
 
 namespace {
 struct BufferizationInlinerInterface : public DialectInlinerInterface {
   using DialectInlinerInterface::DialectInlinerInterface;
 
if (!attr.getValue().isa<AffineMapAttr>()) {
       return op->emitError() << "'" << kBufferLayoutAttrName
                              << "' is expected to be a affine map attribute";
     }
     if (!isa<FunctionOpInterface>(op))
       return op->emitError() << "expected " << attr.getName()
                              << " to be used on function-like operations";
     return success();
   }
+  if (attr.getName() == kEscapeAttrName) {
+    auto arrayAttr = attr.getValue().dyn_cast<ArrayAttr>();
+    if (!arrayAttr)
+      return op->emitError() << "'" << kEscapeAttrName
+                             << "' is expected to be a bool array attribute";
+    if (arrayAttr.size() != op->getNumResults())
+      return op->emitError()
+             << "'" << kEscapeAttrName
+             << "' has wrong number of elements, expected "
+             << op->getNumResults() << ", got " << arrayAttr.size();
+    auto bufferizableOp = dyn_cast<BufferizableOpInterface>(op);
+    if (!bufferizableOp)
+      return op->emitError()
+             << "'" << kEscapeAttrName << "' only valid on bufferizable ops";
+    for (const auto &it : llvm::enumerate(arrayAttr)) {
+      auto attr = it.value();
+      auto boolAttr = attr.dyn_cast<BoolAttr>();
+      if (!boolAttr)
+        return op->emitError() << "'" << kEscapeAttrName
+                               << "' is expected to be a bool array attribute";
+      if (!boolAttr.getValue())
+        continue;
+      if (!op->getResult(it.index()).getType().isa<TensorType>())
+        return op->emitError()
+               << "'" << kEscapeAttrName << "' only valid for tensor results";
+      if (!bufferizableOp.bufferizesToAllocation(op->getOpResult(it.index())))
+        return op->emitError() << "'" << kEscapeAttrName
+                               << "' only valid for allocation results";
+    }
+    return success();
+  }
 
   return op->emitError() << "attribute '" << attr.getName()
                          << "' not supported by the bufferization dialect";
 }

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

 
 //===----------------------------------------------------------------------===//
 // AllocTensorOp
 //===----------------------------------------------------------------------===//
 
 LogicalResult AllocTensorOp::bufferize(RewriterBase &rewriter,
                                        const BufferizationOptions &options) {
   OpBuilder::InsertionGuard g(rewriter);
+  Operation *op = this->getOperation();
   Location loc = getLoc();
 
   // Nothing to do for dead AllocTensorOps.
   if (getOperation()->getUses().empty()) {
     rewriter.eraseOp(getOperation());
     return success();
   }
 
LogicalResult AllocTensorOp::bufferize(RewriterBase &rewriter,
   if (getCopy()) {
     if (failed(options.createMemCpy(rewriter, loc, copyBuffer, *alloc)))
       return failure();
   }
 
   // Should the buffer be deallocated?
   AnalysisState analysisState(options);
   bool dealloc;
-  if (getEscape()) {
-    dealloc = !*getEscape();
+  if (op->hasAttr(BufferizationDialect::kEscapeAttrName)) {
+    // AllocTensorOp has one result.
+    ArrayAttr escapeAttr =
+        op->getAttr(BufferizationDialect::kEscapeAttrName).cast<ArrayAttr>();
+    dealloc = !escapeAttr[0].cast<BoolAttr>().getValue();
   } else {
     // No "escape" annotation found.
     if (options.createDeallocs) {
       // Perform an ad-hoc analysis.
       dealloc = !analysisState.isTensorYielded(getResult());
     } else {
       dealloc = false;
     }
return emitError("expected ")
            << getType().getNumDynamicDims() << " dynamic sizes";
   if (getCopy() && getCopy().getType() != getType())
     return emitError("expected that `copy` and return type match");
   return success();
 }
 
 void AllocTensorOp::build(OpBuilder &builder, OperationState &result,
                           RankedTensorType type, ValueRange dynamicSizes) {
-  build(builder, result, type, dynamicSizes, /*copy=*/Value(),
-        /*escape=*/BoolAttr());
-}
-
-void AllocTensorOp::build(OpBuilder &builder, OperationState &result,
-                          RankedTensorType type, ValueRange dynamicSizes,
-                          Value copy) {
-  build(builder, result, type, dynamicSizes, copy, /*escape=*/BoolAttr());
-}
-
-void AllocTensorOp::build(OpBuilder &builder, OperationState &result,
-                          RankedTensorType type, ValueRange dynamicSizes,
-                          Value copy, bool escape) {
-  build(builder, result, type, dynamicSizes, copy, builder.getBoolAttr(escape));
+  build(builder, result, type, dynamicSizes, /*copy=*/Value());
 }
 
 namespace {
 /// Change the type of the result of a `bufferization.alloc_tensor` by making
 /// the result type statically sized along dimension that in the original
 /// operation where defined as dynamic, but the size was defined using a
 /// `constant` op. For example:
 ///
for (int64_t i = 0; i < op.getType().getRank(); ++i) {
         newDynamicSizes.push_back(value);
       }
     }
     RankedTensorType newType = RankedTensorType::get(
         newShape, op.getType().getElementType(), op.getType().getEncoding());
     if (newType == op.getType())
       return failure();
     auto newOp = rewriter.create<AllocTensorOp>(
-        op.getLoc(), newType, newDynamicSizes, /*copy=*/Value(),
-        /*escape=*/op.getEscapeAttr());
+        op.getLoc(), newType, newDynamicSizes, /*copy=*/Value());
     rewriter.replaceOpWithNewOp<tensor::CastOp>(op, op.getType(), newOp);
     return success();
   }
 };
 
 struct FoldDimOfAllocTensorOp : public OpRewritePattern<tensor::DimOp> {
   using OpRewritePattern<tensor::DimOp>::OpRewritePattern;
 

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

LogicalResult mlir::bufferization::insertTensorCopies(
 
   return insertTensorCopies(op, state);
 }
 
 LogicalResult
 mlir::bufferization::insertTensorCopies(Operation *op,
                                         const AnalysisState &state) {
   IRRewriter rewriter(op->getContext());
+  StringRef escapeAttrName = BufferizationDialect::kEscapeAttrName;
+
   WalkResult result = op->walk([&](Operation *op) {
     auto bufferizableOp = state.getOptions().dynCastBufferizableOp(op);
     if (!bufferizableOp)
       return WalkResult::skip();
 
-    // Find AllocTensorOps without an `escape` attribute and add the attribute
+    // Find allocations without an `escape` attribute and add the attribute
     // based on analysis results.
-    if (auto allocTensorOp = dyn_cast<AllocTensorOp>(op)) {
-      if (allocTensorOp.getEscape())
-        return WalkResult::advance();
+    if (!op->hasAttr(escapeAttrName)) {
+      SmallVector<bool> escapeAttrValue;
+      bool foundTensorResult = false;
+      for (OpResult opResult : op->getOpResults()) {
+        if (!opResult.getType().isa<TensorType>() ||
+            !bufferizableOp.bufferizesToAllocation(opResult)) {
+          escapeAttrValue.push_back(false);
+          continue;
+        }
+        foundTensorResult = true;
         bool escape = !state.getOptions().createDeallocs ||
-                    state.isTensorYielded(allocTensorOp.getResult());
-      allocTensorOp.setEscapeAttr(rewriter.getBoolAttr(escape));
-      return WalkResult::advance();
+                      state.isTensorYielded(opResult);
+        escapeAttrValue.push_back(escape);
+      }
+      if (foundTensorResult)
+        op->setAttr(escapeAttrName, rewriter.getBoolArrayAttr(escapeAttrValue));
     }
 
     // Find inplacability conflicts and resolve them. (Typically with explicit
     // tensor copies in the form of AllocTensorOps.)
     rewriter.setInsertionPoint(op);
     if (failed(bufferizableOp.resolveConflicts(rewriter, state)))
       return WalkResult::interrupt();
 

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

LogicalResult resolveConflicts(Operation *op, RewriterBase &rewriter,
     SmallVector<Value> yieldValues;
     for (int64_t idx = 0;
          idx < static_cast<int64_t>(yieldOp.getResults().size()); ++idx) {
       Value value = yieldOp.getResults()[idx];
       if (!indices.contains(idx) || equivalentYields.contains(idx)) {
         yieldValues.push_back(value);
         continue;
       }
-      Value alloc = rewriter.create<bufferization::AllocTensorOp>(
-          yieldOp.getLoc(), value.getType().cast<RankedTensorType>(),
-          /*dynamicSizes=*/ValueRange(), value, /*escape=*/true);
+      Value alloc = allocateTensorForShapedValue(rewriter, yieldOp.getLoc(),
+                                                 value, /*escape=*/true);
       yieldValues.push_back(alloc);
     }
 
     rewriter.updateRootInPlace(
         yieldOp, [&]() { yieldOp.getResultsMutable().assign(yieldValues); });
     return success();
   }
 
LogicalResult resolveConflicts(Operation *op, RewriterBase &rewriter,
     for (int64_t idx = 0;
          idx < static_cast<int64_t>(conditionOp.getArgs().size()); ++idx) {
       Value value = conditionOp.getArgs()[idx];
       if (!indicesBefore.contains(idx) ||
           equivalentYieldsBefore.contains(idx)) {
         beforeYieldValues.push_back(value);
         continue;
       }
-      Value alloc = rewriter.create<bufferization::AllocTensorOp>(
-          conditionOp.getLoc(), value.getType().cast<RankedTensorType>(),
-          /*dynamicSizes=*/ValueRange(), value, /*escape=*/true);
+      Value alloc = allocateTensorForShapedValue(rewriter, conditionOp.getLoc(),
+                                                 value, /*escape=*/true);
       beforeYieldValues.push_back(alloc);
     }
     rewriter.updateRootInPlace(conditionOp, [&]() {
       conditionOp.getArgsMutable().assign(beforeYieldValues);
     });
 
     // Update "after" region.
     rewriter.setInsertionPoint(yieldOp);
     SmallVector<Value> afterYieldValues;
     for (int64_t idx = 0;
          idx < static_cast<int64_t>(yieldOp.getResults().size()); ++idx) {
       Value value = yieldOp.getResults()[idx];
       if (!indicesAfter.contains(idx) || equivalentYieldsAfter.contains(idx)) {
         afterYieldValues.push_back(value);
         continue;
       }
-      Value alloc = rewriter.create<bufferization::AllocTensorOp>(
-          yieldOp.getLoc(), value.getType().cast<RankedTensorType>(),
-          /*dynamicSizes=*/ValueRange(), value, /*escape=*/true);
+      Value alloc = allocateTensorForShapedValue(rewriter, yieldOp.getLoc(),
+                                                 value, /*escape=*/true);
       afterYieldValues.push_back(alloc);
     }
     rewriter.updateRootInPlace(yieldOp, [&]() {
       yieldOp.getResultsMutable().assign(afterYieldValues);
     });
 
     return success();
   }
for (OpResult opResult : foreachThreadOp->getOpResults()) {
              "expected exactly one aliasing OpOperand");
       assert(isa<ParallelInsertSliceOp>(destOperands.front()->getOwner()) &&
              "expected ParallelInsertSliceOp");
 
       // Nothing to do if there is no conflict.
       if (state.isInPlace(*destOperands.front()))
         continue;
 
-      // Create AllocTensorOp.
+      // Insert tensor allocation.
       bool isYielded = state.isTensorYielded(opResult);
-      auto resultType = opResult.getType().cast<RankedTensorType>();
-      Value alloc = rewriter.create<bufferization::AllocTensorOp>(
-          op->getLoc(), resultType, /*dynamicDims=*/ValueRange(),
-          /*copy=*/destOperands.front()->get(),
+      Value alloc = allocateTensorForShapedValue(rewriter, op->getLoc(),
+                                                 destOperands.front()->get(),
                                                  /*escape=*/isYielded);
 
       // Update terminator operand.
       rewriter.updateRootInPlace(destOperands.front()->getOwner(),
                                  [&]() { destOperands.front()->set(alloc); });
     }
 
     return success();
   }

mlir/test/Dialect/Bufferization/Transforms/tensor-copy-insertion.mlir

 // RUN: mlir-opt %s -tensor-copy-insertion -split-input-file | FileCheck %s
 // RUN: mlir-opt %s -tensor-copy-insertion="bufferize-function-boundaries allow-return-allocs" -split-input-file | FileCheck %s --check-prefix=CHECK-FUNC
 // RUN: mlir-opt %s -tensor-copy-insertion="create-deallocs=0" -split-input-file | FileCheck %s --check-prefix=CHECK-NO-DEALLOC
 
 // CHECK-LABEL: func @read_after_write_conflict(
 //  CHECK-SAME:     %[[t:.*]]: tensor<?xf32>
 // CHECK-FUNC-LABEL: func @read_after_write_conflict(
 // CHECK-NO-DEALLOC-LABEL: func @read_after_write_conflict(
 func.func @read_after_write_conflict(%t: tensor<?xf32>, %idx: index, %f: f32)
   -> (tensor<?xf32>, tensor<?xf32>)
 {
-  // CHECK: %[[copy:.*]] = bufferization.alloc_tensor() copy(%[[t]]) {escape = false} : tensor<?xf32>
-  // CHECK-FUNC: bufferization.alloc_tensor() copy(%{{.*}}) {escape = true} : tensor<?xf32>
-  // CHECK-NO-DEALLOC: bufferization.alloc_tensor() copy(%{{.*}}) {escape = true} : tensor<?xf32>
+  // CHECK: %[[copy:.*]] = bufferization.alloc_tensor() copy(%[[t]]) {bufferization.escape = [false]} : tensor<?xf32>
+  // CHECK-FUNC: bufferization.alloc_tensor() copy(%{{.*}}) {bufferization.escape = [true]} : tensor<?xf32>
+  // CHECK-NO-DEALLOC: bufferization.alloc_tensor() copy(%{{.*}}) {bufferization.escape = [true]} : tensor<?xf32>
   // CHECK: %[[insert:.*]] = tensor.insert %{{.*}} into %[[copy]]
   %0 = tensor.insert %f into %t[%idx] : tensor<?xf32>
   // CHECK: return %[[insert]], %[[t]]
   return %0, %t : tensor<?xf32>, tensor<?xf32>
 }
 
 // -----
 
 // CHECK-LABEL: func @return_alloc_tensor
 // CHECK-FUNC-LABEL: func @return_alloc_tensor
 // CHECK-NO-DEALLOC-LABEL: func @return_alloc_tensor
 func.func @return_alloc_tensor() -> (tensor<5xf32>) {
-  // CHECK: bufferization.alloc_tensor() {escape = false} : tensor<5xf32>
-  // CHECK-FUNC: bufferization.alloc_tensor() {escape = true} : tensor<5xf32>
-  // CHECK-NO-DEALLOC: bufferization.alloc_tensor() {escape = true} : tensor<5xf32>
+  // CHECK: bufferization.alloc_tensor() {bufferization.escape = [false]} : tensor<5xf32>
+  // CHECK-FUNC: bufferization.alloc_tensor() {bufferization.escape = [true]} : tensor<5xf32>
+  // CHECK-NO-DEALLOC: bufferization.alloc_tensor() {bufferization.escape = [true]} : tensor<5xf32>
   %0 = bufferization.alloc_tensor() : tensor<5xf32>
   return %0 : tensor<5xf32>
 }
 
 // -----
 
 // CHECK-LABEL: func @do_not_copy_undefined_tensor
 // CHECK-NO-DEALLOC-LABEL: func @do_not_copy_undefined_tensor
 func.func @do_not_copy_undefined_tensor(%f: f32, %idx: index)
   -> (tensor<5xf32>, tensor<5xf32>)
 {
-  // CHECK: bufferization.alloc_tensor() {escape = false} : tensor<5xf32>
+  // CHECK: bufferization.alloc_tensor() {bufferization.escape = [false]} : tensor<5xf32>
   // The second alloc_tensor should not have a copy operand.
-  // CHECK: bufferization.alloc_tensor() {escape = false} : tensor<5xf32>
+  // CHECK: bufferization.alloc_tensor() {bufferization.escape = [false]} : tensor<5xf32>
 
-  // CHECK-NO-DEALLOC: bufferization.alloc_tensor() {escape = true} : tensor<5xf32>
-  // CHECK-NO-DEALLOC: bufferization.alloc_tensor() {escape = true} : tensor<5xf32>
+  // CHECK-NO-DEALLOC: bufferization.alloc_tensor() {bufferization.escape = [true]} : tensor<5xf32>
+  // CHECK-NO-DEALLOC: bufferization.alloc_tensor() {bufferization.escape = [true]} : tensor<5xf32>
   %0 = bufferization.alloc_tensor() : tensor<5xf32>
   %1 = tensor.insert %f into %0[%idx] : tensor<5xf32>
   return %0, %1 : tensor<5xf32>, tensor<5xf32>
 }
 
 // -----
 
 // CHECK-LABEL: func @do_not_copy_when_overwritten
 func.func @do_not_copy_when_overwritten(%t: tensor<5xf32>, %f: f32)
   -> (tensor<5xf32>, tensor<5xf32>)
 {
-  // CHECK: %[[alloc:.*]] = bufferization.alloc_tensor() {escape = false} : tensor<5xf32>
+  // CHECK: %[[alloc:.*]] = bufferization.alloc_tensor() {bufferization.escape = [false]} : tensor<5xf32>
   // CHECK: linalg.generic {{.*}} outs(%[[alloc]] : tensor<5xf32>)
   %r = linalg.generic {
     indexing_maps = [affine_map<(d0) -> (d0)>],
     iterator_types = ["parallel"]}
     outs(%t : tensor<5xf32>) {
       ^bb0(%arg0 : f32) :
         linalg.yield %f : f32
     } -> tensor<5xf32>
   return %t, %r : tensor<5xf32>, tensor<5xf32>
 }
 
 // -----
 
 // CHECK-LABEL: func @do_not_copy_when_result_not_read
 func.func @do_not_copy_when_result_not_read(%t: tensor<5xf32>, %f: f32)
   -> (tensor<3xf32>)
 {
   %0 = tensor.extract_slice %t[0][3][1] : tensor<5xf32> to tensor<3xf32>
-  // CHECK: %[[alloc:.*]] = bufferization.alloc_tensor() {escape = false} : tensor<3xf32>
+  // CHECK: %[[alloc:.*]] = bufferization.alloc_tensor() {bufferization.escape = [false]} : tensor<3xf32>
   // CHECK: linalg.generic {{.*}} outs(%[[alloc]] : tensor<3xf32>)
   %r = linalg.generic {
     indexing_maps = [affine_map<(d0) -> (d0)>],
     iterator_types = ["parallel"]}
     outs(%0 : tensor<3xf32>) {
       ^bb0(%arg0 : f32) :
         linalg.yield %f : f32
     } -> tensor<3xf32>
   return %r : tensor<3xf32>
 }

mlir/test/Dialect/Bufferization/invalid.mlir

func.func @alloc_tensor_copy_and_dims(%t: tensor<?xf32>, %sz: index) {
   // expected-error @+1{{dynamic sizes not needed when copying a tensor}}
   %0 = bufferization.alloc_tensor(%sz) copy(%t) : tensor<?xf32>
   return
 }
 
 // -----
 
 func.func @alloc_tensor_invalid_escape_attr(%sz: index) {
-  // expected-error @+1{{op attribute 'escape' failed to satisfy constraint: bool attribute}}
-  %0 = bufferization.alloc_tensor(%sz) {escape = 5} : tensor<?xf32>
+  // expected-error @+1{{'bufferization.escape' is expected to be a bool array attribute}}
+  %0 = bufferization.alloc_tensor(%sz) {bufferization.escape = 5} : tensor<?xf32>
   return
 }
+
+// -----
+
+func.func @alloc_tensor_invalid_escape_attr_size(%sz: index) {
+  // expected-error @+1{{'bufferization.escape' has wrong number of elements, expected 1, got 2}}
+  %0 = bufferization.alloc_tensor(%sz) {bufferization.escape = [true, false]} : tensor<?xf32>
+  return
+}
+
+// -----
+
+func.func @escape_attr_non_allocating(%t0: tensor<?xf32>) {
+  // expected-error @+1{{'bufferization.escape' only valid for allocation results}}
+  %0 = tensor.extract_slice %t0[0][5][1] {bufferization.escape = [true]} : tensor<?xf32> to tensor<5xf32>
+  return
+}
+
+// -----
+
+func.func @escape_attr_non_bufferizable(%m0: memref<?xf32>) {
+  // expected-error @+1{{'bufferization.escape' only valid on bufferizable ops}}
+  %0 = memref.cast %m0 {bufferization.escape = [true]} : memref<?xf32> to memref<10xf32>
+  return
+}
 No newline at end of file

mlir/test/Dialect/SCF/one-shot-bufferize-tensor-copy-insertion.mlir

 // RUN: mlir-opt %s -tensor-copy-insertion="allow-return-allocs" -allow-unregistered-dialect -split-input-file | FileCheck %s
 // RUN: mlir-opt %s -tensor-copy-insertion="bufferize-function-boundaries allow-return-allocs" -split-input-file | FileCheck %s --check-prefix=CHECK-FUNC
 
 // CHECK-LABEL: func @scf_for(
 //  CHECK-SAME:     %[[A:.*]]: tensor<?xf32>, %[[B:.*]]: tensor<?xf32>
 func.func @scf_for(%A : tensor<?xf32>, %B : tensor<?xf32>,
                    %lb : index, %ub : index, %step : index)
   -> (tensor<?xf32>, tensor<?xf32>)
 {
-  // CHECK: %[[A_copy:.*]] = bufferization.alloc_tensor() copy(%[[A]]) {escape = false} : tensor<?xf32>
-  // CHECK: %[[B_copy:.*]] = bufferization.alloc_tensor() copy(%[[B]]) {escape = false} : tensor<?xf32>
+  // CHECK: %[[A_copy:.*]] = bufferization.alloc_tensor() copy(%[[A]]) {bufferization.escape = [false]} : tensor<?xf32>
+  // CHECK: %[[B_copy:.*]] = bufferization.alloc_tensor() copy(%[[B]]) {bufferization.escape = [false]} : tensor<?xf32>
   // CHECK:   %[[for:.*]]:2 = scf.for {{.*}} iter_args(%[[iter1:.*]] = %[[A_copy]], %[[iter2:.*]] = %[[B_copy]])
   %r0:2 = scf.for %i = %lb to %ub step %step iter_args(%tA = %A, %tB = %B)
       -> (tensor<?xf32>, tensor<?xf32>)
   {
     // CHECK: scf.yield %[[iter1]], %[[iter2]]
     scf.yield %tA, %tB : tensor<?xf32>, tensor<?xf32>
   }
 
   return %r0#0, %r0#1 : tensor<?xf32>, tensor<?xf32>
 }
 
 // -----
 
 // CHECK-LABEL: func @scf_for_swapping_yields(
 //  CHECK-SAME:     %[[A:.*]]: tensor<?xf32>, %[[B:.*]]: tensor<?xf32>
 func.func @scf_for_swapping_yields(%A : tensor<?xf32>, %B : tensor<?xf32>,
                                    %lb : index, %ub : index, %step : index)
   -> (tensor<?xf32>, tensor<?xf32>)
 {
-  // CHECK: %[[A_copy:.*]] = bufferization.alloc_tensor() copy(%[[A]]) {escape = false} : tensor<?xf32>
-  // CHECK: %[[B_copy:.*]] = bufferization.alloc_tensor() copy(%[[B]]) {escape = false} : tensor<?xf32>
+  // CHECK: %[[A_copy:.*]] = bufferization.alloc_tensor() copy(%[[A]]) {bufferization.escape = [false]} : tensor<?xf32>
+  // CHECK: %[[B_copy:.*]] = bufferization.alloc_tensor() copy(%[[B]]) {bufferization.escape = [false]} : tensor<?xf32>
   // CHECK:   %[[for:.*]]:2 = scf.for {{.*}} iter_args(%[[iter1:.*]] = %[[A_copy]], %[[iter2:.*]] = %[[B_copy]])
   %r0:2 = scf.for %i = %lb to %ub step %step iter_args(%tA = %A, %tB = %B)
       -> (tensor<?xf32>, tensor<?xf32>)
   {
     // Yield tensors in different order.
-    // CHECK-DAG: %[[yield1:.*]] = bufferization.alloc_tensor() copy(%[[iter2]]) {escape = true} : tensor<?xf32>
-    // CHECK-DAG: %[[yield2:.*]] = bufferization.alloc_tensor() copy(%[[iter1]]) {escape = true} : tensor<?xf32>
+    // CHECK-DAG: %[[yield1:.*]] = bufferization.alloc_tensor() copy(%[[iter2]]) {bufferization.escape = [true]} : tensor<?xf32>
+    // CHECK-DAG: %[[yield2:.*]] = bufferization.alloc_tensor() copy(%[[iter1]]) {bufferization.escape = [true]} : tensor<?xf32>
     // CHECK: scf.yield %[[yield1]], %[[yield2]]
     scf.yield %tB, %tA : tensor<?xf32>, tensor<?xf32>
   }
 
   return %r0#0, %r0#1 : tensor<?xf32>, tensor<?xf32>
 }
 
 // -----
 
 // CHECK-LABEL: func @scf_while(
 //  CHECK-SAME:     %[[A:.*]]: tensor<5xi1>, %[[B:.*]]: tensor<5xi1>
 func.func @scf_while(%A: tensor<5xi1>, %B: tensor<5xi1>, %idx: index)
   -> (tensor<5xi1>, tensor<5xi1>)
 {
-  // CHECK: %[[A_copy:.*]] = bufferization.alloc_tensor() copy(%[[A]]) {escape = false} : tensor<5xi1>
-  // CHECK: %[[B_copy:.*]] = bufferization.alloc_tensor() copy(%[[B]]) {escape = false} : tensor<5xi1>
+  // CHECK: %[[A_copy:.*]] = bufferization.alloc_tensor() copy(%[[A]]) {bufferization.escape = [false]} : tensor<5xi1>
+  // CHECK: %[[B_copy:.*]] = bufferization.alloc_tensor() copy(%[[B]]) {bufferization.escape = [false]} : tensor<5xi1>
   // CHECK: %[[loop:.*]]:2 = scf.while (%[[w0:.*]] = %[[A_copy]], %[[w1:.*]] = %[[B_copy]]) {{.*}} {
   %r0, %r1 = scf.while (%w0 = %A, %w1 = %B)
       : (tensor<5xi1>, tensor<5xi1>) -> (tensor<5xi1>, tensor<5xi1>) {
     // CHECK: %[[condition:.*]] = tensor.extract %[[w0]]
     %condition = tensor.extract %w0[%idx] : tensor<5xi1>
     // Yield tensors in different order.
     // CHECK: scf.condition(%[[condition]]) %[[w0]], %[[w1]]
     scf.condition(%condition) %w0, %w1 : tensor<5xi1>, tensor<5xi1>

 
 // CHECK-LABEL: func @scf_while_non_equiv_condition_and_body(
 //  CHECK-SAME:     %[[A:.*]]: tensor<5xi1>, %[[B:.*]]: tensor<5xi1>
 func.func @scf_while_non_equiv_condition_and_body(%A: tensor<5xi1>,
                                                   %B: tensor<5xi1>,
                                                   %idx: index)
   -> (tensor<5xi1>, tensor<5xi1>)
 {
-  // CHECK: %[[A_copy:.*]] = bufferization.alloc_tensor() copy(%[[A]]) {escape = false} : tensor<5xi1>
-  // CHECK: %[[B_copy:.*]] = bufferization.alloc_tensor() copy(%[[B]]) {escape = false} : tensor<5xi1>
+  // CHECK: %[[A_copy:.*]] = bufferization.alloc_tensor() copy(%[[A]]) {bufferization.escape = [false]} : tensor<5xi1>
+  // CHECK: %[[B_copy:.*]] = bufferization.alloc_tensor() copy(%[[B]]) {bufferization.escape = [false]} : tensor<5xi1>
   // CHECK: %[[loop:.*]]:2 = scf.while (%[[w0:.*]] = %[[A_copy]], %[[w1:.*]] = %[[B_copy]]) {{.*}} {
   %r0, %r1 = scf.while (%w0 = %A, %w1 = %B)
       : (tensor<5xi1>, tensor<5xi1>) -> (tensor<5xi1>, tensor<5xi1>) {
     // CHECK: %[[condition:.*]] = tensor.extract %[[w0]]
     %condition = tensor.extract %w0[%idx] : tensor<5xi1>
     // Yield tensors in different order.
-    // CHECK-DAG: %[[yield0:.*]] = bufferization.alloc_tensor() copy(%[[w1]]) {escape = true} : tensor<5xi1>
-    // CHECK-DAG: %[[yield1:.*]] = bufferization.alloc_tensor() copy(%[[w0]]) {escape = true} : tensor<5xi1>
+    // CHECK-DAG: %[[yield0:.*]] = bufferization.alloc_tensor() copy(%[[w1]]) {bufferization.escape = [true]} : tensor<5xi1>
+    // CHECK-DAG: %[[yield1:.*]] = bufferization.alloc_tensor() copy(%[[w0]]) {bufferization.escape = [true]} : tensor<5xi1>
     // CHECK: scf.condition(%[[condition]]) %[[yield0]], %[[yield1]]
     scf.condition(%condition) %w1, %w0 : tensor<5xi1>, tensor<5xi1>
   } do {
   ^bb0(%b0: tensor<5xi1>, %b1: tensor<5xi1>):
     // CHECK: } do {
     // CHECK: ^bb0(%[[b0:.*]]: tensor<5xi1>, %[[b1:.*]]: tensor<5xi1>):
-    // CHECK-DAG: %[[yield2:.*]] = bufferization.alloc_tensor() copy(%[[b1]]) {escape = true} : tensor<5xi1>
-    // CHECK-DAG: %[[yield3:.*]] = bufferization.alloc_tensor() copy(%[[b0]]) {escape = true} : tensor<5xi1>
+    // CHECK-DAG: %[[yield2:.*]] = bufferization.alloc_tensor() copy(%[[b1]]) {bufferization.escape = [true]} : tensor<5xi1>
+    // CHECK-DAG: %[[yield3:.*]] = bufferization.alloc_tensor() copy(%[[b0]]) {bufferization.escape = [true]} : tensor<5xi1>
     // CHECK: scf.yield %[[yield2]], %[[yield3]]
     // CHECK: }
     scf.yield %b1, %b0 : tensor<5xi1>, tensor<5xi1>
   }
 
   return %r0, %r1 : tensor<5xi1>, tensor<5xi1>
 }
 
 // -----
 
 // CHECK-LABEL: func @scf_foreach_thread_out_of_place(
 //  CHECK-SAME:     %[[arg0:.*]]: tensor<100xf32>, %[[arg1:.*]]: tensor<100xf32>
 // CHECK-FUNC-LABEL: func @scf_foreach_thread_out_of_place(
 func.func @scf_foreach_thread_out_of_place(%in: tensor<100xf32>,
                                            %out: tensor<100xf32>) {
   %c1 = arith.constant 1 : index
   %num_threads = arith.constant 100 : index
 
   // CHECK-FUNC-NOT: alloc_tensor
-  // CHECK: %[[alloc:.*]] = bufferization.alloc_tensor() copy(%[[arg1]]) {escape = false} : tensor<100xf32>
+  // CHECK: %[[alloc:.*]] = bufferization.alloc_tensor() copy(%[[arg1]]) {bufferization.escape = [false]} : tensor<100xf32>
   // CHECK: scf.foreach_thread
   %result = scf.foreach_thread (%thread_idx) in (%num_threads) -> tensor<100xf32> {
       // CHECK: tensor.extract_slice
       // CHECK: scf.foreach_thread.perform_concurrently
       // CHECK: scf.foreach_thread.parallel_insert_slice %{{.*}} into %[[alloc]]
       %1 = tensor.extract_slice %in[%thread_idx][1][1] : tensor<100xf32> to tensor<1xf32>
       scf.foreach_thread.perform_concurrently {
         scf.foreach_thread.parallel_insert_slice %1 into %out[%thread_idx][1][1] :
           tensor<1xf32> into tensor<100xf32>
       }
   }
   return
 }

mlir/test/Dialect/Tensor/one-shot-bufferize-tensor-copy-insertion.mlir

 // RUN: mlir-opt %s -tensor-copy-insertion -split-input-file | FileCheck %s
 // RUN: mlir-opt %s -tensor-copy-insertion="bufferize-function-boundaries allow-return-allocs" -split-input-file | FileCheck %s --check-prefix=CHECK-FUNC
 
 // CHECK-LABEL: func @extract_slice(
 //  CHECK-SAME:     %[[t:.*]]: tensor<?xf32>
 // CHECK-FUNC-LABEL: func @extract_slice(
 func.func @extract_slice(%t: tensor<?xf32>, %idx: index, %f: f32)
   -> (tensor<5xf32>, tensor<?xf32>)
 {
   // CHECK: %[[extract_slice:.*]] = tensor.extract_slice %[[t]][10] [5] [1]
   %0 = tensor.extract_slice %t[10][5][1] : tensor<?xf32> to tensor<5xf32>
-  // CHECK: %[[alloc:.*]] = bufferization.alloc_tensor() copy(%[[extract_slice]]) {escape = false} : tensor<5xf32>
-  // CHECK-FUNC: bufferization.alloc_tensor() copy(%{{.*}}) {escape = true} : tensor<5xf32>
+  // CHECK: %[[alloc:.*]] = bufferization.alloc_tensor() copy(%[[extract_slice]]) {bufferization.escape = [false]} : tensor<5xf32>
+  // CHECK-FUNC: bufferization.alloc_tensor() copy(%{{.*}}) {bufferization.escape = [true]} : tensor<5xf32>
   // CHECK: %[[insert:.*]] = tensor.insert %{{.*}} into %[[alloc]]
   %1 = tensor.insert %f into %0[%idx] : tensor<5xf32>
   // CHECK: return %[[insert]], %[[t]]
   return %1, %t : tensor<5xf32>, tensor<?xf32>
 }