[mlir][bufferize] Add memory space inference

[springerm]

Summary

Bufferization usually processes ops top-to-bottom during the rewriting phase. In that case, the memref type (and thus memory space) of an op can be directly inferred from all previous ops. However, there are a few exceptions such as scf.if.

These ops have a body with a terminator that must be bufferized. The terminator is bufferized together with the parent op. In that case, the bufferization infrastructure must compute a memref type and cannot directly take the memref type from a previous (already bufferized) op. E.g.:

scf.if %c -> tensor<?xf32> {
  // ...
  // To bufferize this op, the memory space of buffer(%r) is needed.
  scf.yield %r : tensor<?xf32>
}

In the above example, the bufferized type of %r is unknown when bufferizing the scf.yield op.

This change adds an ad-hoc analysis that inspects reverse SSA use-def chains of tensor values until a tensor is found for which the memory space is known.

Note: This change sets up the infrastructure and updates the BufferizableOpInterface implementations of the "arith" dialect to provide a meaningful test case. The BufferizableOpInterface implementations of other dialects are updated in subsequent changes.

Depends On D128278

[springerm]

update tests

[springerm]

major refactoring and rolling over changes into a different revision

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

 
 /// Return the buffer type for a given Value (tensor) after bufferization.
 ///
 /// Note: Op implementations should preferrably call `getBuffer()->getType()`.
 /// This function should only be used if `getBuffer` cannot be used.
 FailureOr<BaseMemRefType> getBufferType(Value value,
                                         const BufferizationOptions &options);
 
+/// Try to infer the memory space from the given tensor value.
+FailureOr<unsigned> inferMemorySpace(Value value,
+                                     const BufferizationOptions &options);
+
+/// Try to infer the memory space from the given tensor value. If no memory
+/// space could be inferred, take the default memory space from `options`. If no
+/// default memory space is specified, return `failure`.
+FailureOr<unsigned>
+inferOrDefaultMemorySpace(Value value, const BufferizationOptions &options);
+
+/// Try to infer the memory space from the given tensor value.
+FailureOr<unsigned> inferMemorySpace(OpResult opResult,
+                                     const AnalysisState &state);
+
 /// Replace an op with replacement values. The op is deleted. Tensor OpResults
 /// must be replaced with memref values.
 void replaceOpWithBufferizedValues(RewriterBase &rewriter, Operation *op,
                                    ValueRange values);
 
 /// Replace an op with a new op. The new op must have the same number of
 /// results as the replaced op. The new op may not return any tensor values.
 template <typename OpTy, typename... Args>

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

let methods = [
         /*methodBody=*/"",
         /*defaultImplementation=*/[{
           llvm_unreachable("bufferize not implemented");
           return failure();
         }]
       >,
       InterfaceMethod<
         /*desc=*/[{
+          Try to infer the memory space of the given value after bufferization.
+          The value must belong to this op, i.e., it is an OpResult of this op
+          or a BlockArgument of a block of this op.
+
+          Ops that bufferize to an allocation allow users to specify a memory
+          space directly on the op. Such ops should return that memory space if
+          specified.
+
+          By default, this method follows the reverse SSA use-def chains, until
+          a value is found for which a memory space can be inferred. Only values
+          that are guaranteed to bufferize to equivalent buffers are considered
+          when following the chains. Following all potentially aliasing
+          OpOperands is not safe because there could be false positives, i.e.,
+          OpOperands that are not actually aliasing with the given OpResult at
+          runtime; aliasing is a may-be property in BufferizableOpInterface,
+          equivalence is a must-be property.
+
+          Ops that do not necessarily bufferize to equivalent buffers but can
+          provide memory space guarantees with respect to their OpResults or
+          BlockArguments should override this method for better memory space
+          inference.
+        }],
+        /*retType=*/"FailureOr<unsigned>",
+        /*methodName=*/"inferMemorySpace",
+        /*args=*/(ins "Value":$value, "const AnalysisState &":$state),
+        /*methodBody=*/"",
+        /*defaultImplementation=*/[{
+          if (auto bbArg = value.dyn_cast<BlockArgument>()) {
+            assert(bbArg.getOwner()->getParentOp() == $_op.getOperation() &&
+                   "block argument must belong to this op");
+            return failure();
+          }
+          OpResult opResult = value.cast<OpResult>();
+          assert(opResult.getDefiningOp() == $_op.getOperation() &&
+                 "op result must belong to this op");
+          return bufferization::inferMemorySpace(opResult, state);
+        }]
+      >,
+      InterfaceMethod<
+        /*desc=*/[{
           Return `true` if the given Value can be written to in-place. Value is
           either an OpResult of this operation or a BlockArgument of a block of
           this operation.
 
           Most OpResult buffers can be written to, but some ops such as
           ConstantOp may bufferize to non-writable (read-only) memory locations.
           Therefore, by default, this method returns `true` for OpResults. This
           method will never be called on OpResults that do not have a tensor

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

bool bufferizesToMemoryRead(OpOperand &opOperand,
                                 const AnalysisState &state);
 
     bool bufferizesToMemoryWrite(OpOperand &opOperand,
                                  const AnalysisState &state);
 
     SmallVector<OpResult> getAliasingOpResult(
         OpOperand &opOperand, const AnalysisState &state);
 
+    FailureOr<unsigned> inferMemorySpace(Value value,
+                                         const AnalysisState &state);
+
     RankedTensorType getType() {
       return getResult().getType().cast<RankedTensorType>();
     }
 
     // Return true if the size of the tensor is dynamic at `idx`
     bool isDynamicDim(unsigned idx) {
       return getType().isDynamicDim(idx);
     }
LogicalResult bufferize(RewriterBase &rewriter,
       // will start using to_tensor's memref operand. Once all users of
       // to_tensor are bufferized, the op will not have any users anymore and
       // DCE away. In case of partial bufferization, to_memref(to_tensor(x))
       // constructs may be left over. These are folded by the canonicalizer or
       // FinalizingBufferize.
       return success();
     }
 
+    FailureOr<unsigned> inferMemorySpace(Value value,
+                                         const AnalysisState &state) {
+      return getMemref().getType().cast<BaseMemRefType>().getMemorySpaceAsInt();
+    }
+
     bool isWritable(Value value, const AnalysisState &state) const {
       // It is unknown whether the memref operand is writable or not.
       return false;
     }
   }];
 
   let assemblyFormat = "$memref attr-dict `:` type($memref)";
 

mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td

let options = [
     Option<"allowReturnAllocs", "allow-return-allocs", "bool",
             /*default=*/"false",
            "Allows returning/yielding new allocations from a block.">,
     Option<"bufferizeFunctionBoundaries", "bufferize-function-boundaries",
            "bool", /*default=*/"0",
            "Bufferize function boundaries (experimental).">,
     Option<"createDeallocs", "create-deallocs", "bool", /*default=*/"true",
            "Specify if new allocations should be deallocated.">,
+    Option<"mustInferMemorySpace", "must-infer-memory-space", "bool",
+           /*default=*/"false",
+           "The memory space of an memref types must always be inferred. If "
+           "unset, a default memory space of 0 is used otherwise.">,
   ];
   let constructor = "mlir::bufferization::createTensorCopyInsertionPass()";
 }
 
 def AllocTensorElimination : Pass<"eliminate-alloc-tensors"> {
   let summary = "Try to eliminate all alloc_tensor ops.";
   let description = [{
     This pass tries to eliminate all insert_slice op-anchored alloc_tensor ops.

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

struct SelectOpInterface
 
   SmallVector<OpOperand *>
   getAliasingOpOperand(Operation *op, OpResult opResult,
                        const AnalysisState &state) const {
     return {&op->getOpOperand(1) /*true_value*/,
             &op->getOpOperand(2) /*false_value*/};
   }
 
+  FailureOr<unsigned> inferMemorySpace(Operation *op, Value value,
+                                       const AnalysisState &state) const {
+    const BufferizationOptions &options = state.getOptions();
+    auto selectOp = cast<arith::SelectOp>(op);
+#ifndef NDEBUG
+    OpResult opResult = value.cast<OpResult>();
+    assert(opResult.getOwner() == op && "value must be an OpResult of op");
+#endif // NDEBUG
+
+    FailureOr<unsigned> trueMemorySpace =
+        bufferization::inferMemorySpace(selectOp.getTrueValue(), options);
+    FailureOr<unsigned> falseMemorySpace =
+        bufferization::inferMemorySpace(selectOp.getFalseValue(), options);
+
+    if (failed(trueMemorySpace) || failed(falseMemorySpace) ||
+        *trueMemorySpace != *falseMemorySpace)
+      return failure();
+    return *trueMemorySpace;
+  }
+
   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
     // instead of its OpOperands. In the worst case, 2 copies are inserted at
     // the moment (one for each tensor). When copying the op result, only one
     // copy would be needed.
     FailureOr<Value> maybeTrueBuffer =
         getBuffer(rewriter, selectOp.getTrueValue(), options);
     FailureOr<Value> maybeFalseBuffer =
         getBuffer(rewriter, selectOp.getFalseValue(), options);
     if (failed(maybeTrueBuffer) || failed(maybeFalseBuffer))
       return failure();
     Value trueBuffer = *maybeTrueBuffer;
     Value falseBuffer = *maybeFalseBuffer;
 
+    // Infer memory space.
+    // Note: We try to infer the memory space here instead of just taking the
+    // memory space of either true/false operand, so that a better error message
+    // is produced.
+    FailureOr<unsigned> memorySpace =
+        inferOrDefaultMemorySpace(selectOp.getResult(), options);
+    if (failed(memorySpace))
+      return selectOp->emitError("could not infer memory space");
+
     // The "true" and the "false" operands must have the same type. If the
     // buffers have different types, they differ only in their layout map. Cast
     // both of them to the most dynamic MemRef type.
     if (trueBuffer.getType() != falseBuffer.getType()) {
       auto trueType = trueBuffer.getType().cast<MemRefType>();
       int64_t dynamicOffset = ShapedType::kDynamicStrideOrOffset;
       SmallVector<int64_t> dynamicStrides(trueType.getRank(),
                                           ShapedType::kDynamicStrideOrOffset);
       AffineMap stridedLayout = makeStridedLinearLayoutMap(
           dynamicStrides, dynamicOffset, op->getContext());
       auto castedType =
           MemRefType::get(trueType.getShape(), trueType.getElementType(),
-                          stridedLayout, trueType.getMemorySpaceAsInt());
+                          stridedLayout, *memorySpace);
       trueBuffer = rewriter.create<memref::CastOp>(loc, castedType, trueBuffer);
       falseBuffer =
           rewriter.create<memref::CastOp>(loc, castedType, falseBuffer);
     }
 
     replaceOpWithNewBufferizedOp<arith::SelectOp>(
         rewriter, op, selectOp.getCondition(), trueBuffer, falseBuffer);
     return success();

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);
   }
 
+  // Create AllocTensorOp.
   auto allocTensorOp = b.create<AllocTensorOp>(loc, tensorType, dynamicSizes,
                                                copy ? tensor : Value());
   allocTensorOp->setAttr(BufferizationDialect::kEscapeAttrName,
                          b.getBoolArrayAttr({escape}));
+
+  // Add 'memory_space' attribute. Not needed if 'copy' operand is specified.
+  if (copy)
+    return allocTensorOp.getResult();
+  FailureOr<unsigned> maybeMemorySpace =
+      inferOrDefaultMemorySpace(tensor, options);
+  if (failed(maybeMemorySpace)) {
+    if (OpResult opResult = tensor.dyn_cast<OpResult>())
+      return opResult.getDefiningOp()->emitError(
+          "could not infer memory space");
+    return tensor.cast<BlockArgument>()
+        .getParentBlock()
+        ->getParentOp()
+        ->emitError("could not infer memory space");
+  }
+  allocTensorOp->setAttr(BufferizationDialect::kMemorySpaceAttrName,
+                         b.getI64ArrayAttr({*maybeMemorySpace}));
   return allocTensorOp.getResult();
 }
 
 LogicalResult BufferizableOpInterface::resolveTensorOpOperandConflicts(
     RewriterBase &rewriter, const AnalysisState &state) {
   OpBuilder::InsertionGuard g(rewriter);
   Operation *op = getOperation();
   SmallVector<OpOperand *> outOfPlaceOpOperands;
#ifndef NDEBUG
   assert((!rankedTensorType || memrefType.cast<MemRefType>().getRank() ==
                                    rankedTensorType.getRank()) &&
          "to_memref would be invalid: mismatching ranks");
 #endif
 }
 
 FailureOr<Value> bufferization::getBuffer(RewriterBase &rewriter, Value value,
                                           const BufferizationOptions &options) {
+#ifndef NDEBUG
   auto tensorType = value.getType().dyn_cast<TensorType>();
   assert(tensorType && "unexpected non-tensor type");
+#endif // NDEBUG
 
   // Replace "%t = to_tensor %m" with %m.
   if (auto toTensorOp = value.getDefiningOp<bufferization::ToTensorOp>())
     return toTensorOp.getMemref();
 
   // Insert to_memref op.
+  FailureOr<Type> memrefType = getBufferType(value, options);
+  if (failed(memrefType))
+    return failure();
+  ensureToMemrefOpIsValid(value, *memrefType);
   OpBuilder::InsertionGuard g(rewriter);
   setInsertionPointAfter(rewriter, value);
-  Type memrefType = getMemRefType(tensorType, options);
-  ensureToMemrefOpIsValid(value, memrefType);
   return rewriter
-      .create<bufferization::ToMemrefOp>(value.getLoc(), memrefType, value)
+      .create<bufferization::ToMemrefOp>(value.getLoc(), *memrefType, value)
       .getResult();
 }
 
 /// Return the buffer type for a given Value (tensor) after bufferization.
 FailureOr<BaseMemRefType>
 bufferization::getBufferType(Value value, const BufferizationOptions &options) {
   auto tensorType = value.getType().dyn_cast<TensorType>();
   assert(tensorType && "unexpected non-tensor type");
 
   if (auto toTensorOp = value.getDefiningOp<bufferization::ToTensorOp>())
     return toTensorOp.getMemref().getType().cast<BaseMemRefType>();
 
-  return getMemRefType(tensorType, options);
+  FailureOr<unsigned> memorySpace = inferOrDefaultMemorySpace(value, options);
+  if (failed(memorySpace)) {
+    Operation *op =
+        value.isa<OpResult>()
+            ? value.getDefiningOp()
+            : value.cast<BlockArgument>().getParentBlock()->getParentOp();
+    return op->emitError("could not infer memory space");
+  }
+  return getMemRefType(tensorType, options, /*layout=*/None, *memorySpace);
+}
+
+FailureOr<unsigned>
+bufferization::inferMemorySpace(Value value,
+                                const BufferizationOptions &options) {
+  auto bufferizableOp = options.dynCastBufferizableOp(value);
+  if (!bufferizableOp)
+    return failure();
+
+  AnalysisState state(options);
+  return bufferizableOp.inferMemorySpace(value, state);
+}
+
+FailureOr<unsigned>
+bufferization::inferOrDefaultMemorySpace(Value value,
+                                         const BufferizationOptions &options) {
+  FailureOr<unsigned> inferred = inferMemorySpace(value, options);
+  if (succeeded(inferred))
+    return inferred;
+
+  if (options.defaultMemorySpace.hasValue())
+    return *options.defaultMemorySpace;
+
+  return failure();
+}
+
+// Follow reverse SSA use-def chains (only taking into account equivalent
+// tensors), until an op is found for which the memory space is known.
+FailureOr<unsigned>
+bufferization::inferMemorySpace(OpResult opResult, const AnalysisState &state) {
+  auto bufferizableOp =
+      state.getOptions().dynCastBufferizableOp(opResult.getDefiningOp());
+  if (!bufferizableOp)
+    return failure();
+
+  // Query aliasing tensor operands.
+  SmallVector<OpOperand *> alasingOpOperands =
+      bufferizableOp.getAliasingOpOperand(opResult, state);
+  if (alasingOpOperands.empty())
+    return failure();
+  // Only consider equivalent buffers. "Aliasing" is a may-be property in
+  // the BufferizableOpInterface and we cannot have any false negatives here.
+  if (bufferizableOp.bufferRelation(opResult, state) !=
+      BufferRelation::Equivalent)
+    return failure();
+
+  for (OpOperand *opOperand : alasingOpOperands) {
+    Value value = opOperand->get();
+    auto nextBufferizableOp = state.getOptions().dynCastBufferizableOp(value);
+    if (!nextBufferizableOp)
+      continue;
+    FailureOr<unsigned> nextResult =
+        nextBufferizableOp.inferMemorySpace(value, state);
+    if (succeeded(nextResult))
+      return *nextResult;
+  }
+
+  return failure();
 }
 
 void bufferization::replaceOpWithBufferizedValues(RewriterBase &rewriter,
                                                   Operation *op,
                                                   ValueRange values) {
   assert(values.size() == op->getNumResults() &&
          "expected one value per OpResult");
   OpBuilder::InsertionGuard g(rewriter);

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

LogicalResult AllocTensorOp::bufferize(RewriterBase &rewriter,
   if (getCopy()) {
     FailureOr<Value> maybeCopyBuffer = getBuffer(rewriter, getCopy(), options);
     if (failed(maybeCopyBuffer))
       return failure();
     copyBuffer = *maybeCopyBuffer;
   }
 
   // Compute memory space of this allocation.
-  unsigned memorySpace;
-  if (op->hasAttr(BufferizationDialect::kMemorySpaceAttrName)) {
-    memorySpace = op->getAttrOfType<ArrayAttr>(
-                        BufferizationDialect::kMemorySpaceAttrName)[0]
-                      .cast<IntegerAttr>()
-                      .getValue()
-                      .getZExtValue();
-  } else if (options.defaultMemorySpace.hasValue()) {
-    memorySpace = *options.defaultMemorySpace;
-  } else {
-    return op->emitError("could not infer memory space");
-  }
+  FailureOr<unsigned> inferredMemSpace =
+      inferOrDefaultMemorySpace(getResult(), options);
+  if (failed(inferredMemSpace))
+    return emitError("could not infer memory space");
 
   // Create memory allocation.
   auto allocType =
       MemRefType::get(getType().getShape(), getType().getElementType(),
-                      AffineMap(), memorySpace);
+                      AffineMap(), *inferredMemSpace);
   SmallVector<Value> dynamicDims = getDynamicSizes();
   if (getCopy()) {
     assert(dynamicDims.empty() && "expected either `copy` or `dynamicDims`");
     populateDynamicDimSizes(rewriter, loc, copyBuffer, dynamicDims);
   }
   FailureOr<Value> alloc =
       options.createAlloc(rewriter, loc, allocType, dynamicDims);
   if (failed(alloc))
     return failure();
 
   // Create memory copy (if any).
   if (getCopy()) {
     if (failed(options.createMemCpy(rewriter, loc, copyBuffer, *alloc)))
       return failure();
   }
 
   // Should the buffer be deallocated?
-  AnalysisState analysisState(options);
   bool dealloc;
   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.
+      AnalysisState analysisState(options);
       dealloc = !analysisState.isTensorYielded(getResult());
     } else {
       dealloc = false;
     }
   }
 
   // Replace op.
   replaceOpWithBufferizedValues(rewriter, getOperation(), *alloc);
 
   // Create buffer deallocation (if requested).
   if (!dealloc)
     return success();
 
   rewriter.setInsertionPoint(rewriter.getInsertionBlock()->getTerminator());
   if (failed(options.createDealloc(rewriter, loc, *alloc)))
     return failure();
   return success();
 }
 
+FailureOr<unsigned>
+AllocTensorOp::inferMemorySpace(Value value, const AnalysisState &state) {
+  Operation *op = getOperation();
+  if (op->hasAttr(BufferizationDialect::kMemorySpaceAttrName))
+    // Memory space was specified as an attribute.
+    return op
+        ->getAttrOfType<ArrayAttr>(
+            BufferizationDialect::kMemorySpaceAttrName)[0]
+        .cast<IntegerAttr>()
+        .getValue()
+        .getZExtValue();
+  if (getCopy()) {
+    // Memory space can be inferred from the "copy" operand. But only if the op
+    // is bufferizable.
+    auto bufferizableOp = state.getOptions().dynCastBufferizableOp(getCopy());
+    if (!bufferizableOp)
+      return failure();
+    return bufferizableOp.inferMemorySpace(getCopy(), state);
+  }
+
+  return failure();
+}
+
 bool AllocTensorOp::isMemoryWrite(OpResult opResult,
                                   const AnalysisState &state) {
   // AllocTensorOps do not write unless they have a `copy` value.
   return static_cast<bool>(getCopy());
 }
 
 bool AllocTensorOp::bufferizesToMemoryRead(OpOperand &opOperand,
                                            const AnalysisState &state) {

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

void runOnOperation() override {
     if (options.hasValue()) {
       if (failed(insertTensorCopies(getOperation(), *options)))
         signalPassFailure();
     } else {
       OneShotBufferizationOptions options;
       options.allowReturnAllocs = allowReturnAllocs;
       options.bufferizeFunctionBoundaries = bufferizeFunctionBoundaries;
       options.createDeallocs = createDeallocs;
+      if (mustInferMemorySpace)
+        options.defaultMemorySpace = None;
       if (failed(insertTensorCopies(getOperation(), options)))
         signalPassFailure();
     }
   }
 
 private:
   Optional<OneShotBufferizationOptions> options;
 };

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

SmallVector<OpResult> getAliasingOpResult(Operation *op, OpOperand &opOperand,
     return {};
   }
 
   BufferRelation bufferRelation(Operation *op, OpResult opResult,
                                 const AnalysisState &state) const {
     return BufferRelation::None;
   }
 
+  FailureOr<unsigned> inferMemorySpace(Operation *op, Value value,
+                                       const AnalysisState &state) const {
+    auto extractSliceOp = cast<tensor::ExtractSliceOp>(op);
+    FailureOr<BaseMemRefType> bufferType =
+        getBufferType(extractSliceOp.source(), state.getOptions());
+    if (failed(bufferType))
+      return failure();
+    return bufferType->getMemorySpaceAsInt();
+  }
+
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
                           const BufferizationOptions &options) const {
     auto extractSliceOp = cast<tensor::ExtractSliceOp>(op);
     Location loc = extractSliceOp.getLoc();
 
     // Even if this op was decided to bufferize out-of-place, do not insert the
     // buffer copy yet. This is done later in this function.
     FailureOr<Value> srcMemref =
LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
     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 resultTensorType = reshapeOp.getResult().getType().cast<TensorType>();
-    auto resultMemRefType = getMemRefType(resultTensorType, options);
+    auto resultMemRefType = getMemRefType(
+        resultTensorType, options, /*layout=*/None,
+        srcBuffer->getType().cast<BaseMemRefType>().getMemorySpaceAsInt());
     replaceOpWithNewBufferizedOp<memref::ReshapeOp>(
         rewriter, op, resultMemRefType, *srcBuffer, *shapeBuffer);
     return success();
   }
 };
 
 } // namespace
 } // namespace tensor

mlir/test/Dialect/Arithmetic/one-shot-bufferize-memory-space-invalid.mlir

+// RUN: mlir-opt %s -one-shot-bufferize="must-infer-memory-space" -split-input-file -verify-diagnostics
+
+func.func @inconsistent_memory_space_arith_select(%c: i1) -> tensor<10xf32> {
+  // Selecting tensors with different memory spaces. Such IR cannot be
+  // bufferized.
+  %0 = bufferization.alloc_tensor() {bufferization.memory_space = [0]} : tensor<10xf32>
+  %1 = bufferization.alloc_tensor() {bufferization.memory_space = [1]} : tensor<10xf32>
+  // expected-error @+2 {{could not infer memory space}}
+  // expected-error @+1 {{failed to bufferize op}}
+  %r = arith.select %c, %0, %1 : tensor<10xf32>
+  func.return %r : tensor<10xf32>
+}

mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize-memory-space-invalid.mlir

 // RUN: mlir-opt %s -allow-unregistered-dialect -one-shot-bufferize="must-infer-memory-space" -split-input-file -verify-diagnostics
 
 func.func @alloc_tensor_without_memory_space() -> tensor<10xf32> {
   // expected-error @+2 {{could not infer memory space}}
   // expected-error @+1 {{failed to bufferize op}}
   %0 = bufferization.alloc_tensor() : tensor<10xf32>
   return %0 : tensor<10xf32>
 }
+
+// -----
+
+func.func @memory_space_of_unknown_op() -> f32 {
+  %c0 = arith.constant 0 : index
+  // expected-error @+1 {{could not infer memory space}}
+  %t = "test.dummy_op"() : () -> (tensor<10xf32>)
+  // expected-error @+1 {{failed to bufferize op}}
+  %s = tensor.extract %t[%c0] : tensor<10xf32>
+  return %s : f32
+}

mlir/test/Dialect/Bufferization/Transforms/tensor-copy-insertion-memory-space-invalid.mlir

+// RUN: mlir-opt %s -tensor-copy-insertion="must-infer-memory-space" -split-input-file -verify-diagnostics
+
+// An alloc is inserted but the copy is emitted. Therefore, the memory space
+// should be specified on the alloc_tensor op.
+func.func @memory_space_of_unknown_op() -> (tensor<10xf32>, tensor<10xf32>) {
+  %c0 = arith.constant 0 : index
+  %cst = arith.constant 0.0 : f32
+  // expected-error @+1 {{could not infer memory space}}
+  %t = bufferization.alloc_tensor() : tensor<10xf32>
+  %s = tensor.insert %cst into %t[%c0] : tensor<10xf32>
+  return %s, %t : tensor<10xf32>, tensor<10xf32>
+}

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

+// RUN: mlir-opt %s -tensor-copy-insertion="must-infer-memory-space" -split-input-file | FileCheck %s
+
+// CHECK-LABEL: func @unknown_op_copy
+func.func @unknown_op_copy() -> (tensor<10xf32>, tensor<10xf32>) {
+  %c0 = arith.constant 0 : index
+  %cst = arith.constant 0.0 : f32
+  // CHECK: %[[dummy:.*]] = "test.dummy_op"() : () -> tensor<10xf32>
+  %t = "test.dummy_op"() : () -> tensor<10xf32>
+  // CHECK: %[[copy:.*]] = bufferization.alloc_tensor() copy(%[[dummy]]) {bufferization.escape = [false]} : tensor<10xf32>
+  %s = tensor.insert %cst into %t[%c0] : tensor<10xf32>
+  return %s, %t : tensor<10xf32>, tensor<10xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func @alloc_tensor_copy
+func.func @alloc_tensor_copy() -> (tensor<10xf32>, tensor<10xf32>) {
+  %c0 = arith.constant 0 : index
+  %cst = arith.constant 0.0 : f32
+  // CHECK: bufferization.alloc_tensor() {bufferization.escape = [false], bufferization.memory_space = [1]} : tensor<10xf32>
+  %t = bufferization.alloc_tensor() {bufferization.memory_space = [1]} : tensor<10xf32>
+  // CHECK: bufferization.alloc_tensor() {bufferization.escape = [false], bufferization.memory_space = [1]} : tensor<10xf32>
+  %s = tensor.insert %cst into %t[%c0] : tensor<10xf32>
+  return %s, %t : tensor<10xf32>, tensor<10xf32>
+}

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

 
 // 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() {bufferization.escape = [false]} : tensor<5xf32>
   // The second alloc_tensor should not have a copy operand.
-  // CHECK: bufferization.alloc_tensor() {bufferization.escape = [false]} : tensor<5xf32>
+  // CHECK: bufferization.alloc_tensor() {bufferization.escape = [false], bufferization.memory_space = [0]} : tensor<5xf32>
 
   // CHECK-NO-DEALLOC: bufferization.alloc_tensor() {bufferization.escape = [true]} : tensor<5xf32>
-  // CHECK-NO-DEALLOC: bufferization.alloc_tensor() {bufferization.escape = [true]} : tensor<5xf32>
+  // CHECK-NO-DEALLOC: bufferization.alloc_tensor() {bufferization.escape = [true], bufferization.memory_space = [0]} : 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() {bufferization.escape = [false]} : tensor<5xf32>
+  // CHECK: %[[alloc:.*]] = bufferization.alloc_tensor() {bufferization.escape = [false], bufferization.memory_space = [0]} : 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() {bufferization.escape = [false]} : tensor<3xf32>
+  // CHECK: %[[alloc:.*]] = bufferization.alloc_tensor() {bufferization.escape = [false], bufferization.memory_space = [0]} : 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>
 }