[mlir][bufferization] Make memory_space a dialect attribute

[springerm]

Summary

Some ops are not bufferized directly but rewritten into a combination of other (transitory) ops that can then bufferize independently (e.g., tensor.pad). This change makes it possible to attach memory space information to transitory ops.

Similar to bufferization.escape, bufferization.memory_space is now also an array attribute: One IntegerAttr per OpResult.

[aartbik]

LGTM for sparse (since mechanical)

[springerm]

After discussing this offline, there is probably a better solution: AllocationScopeOpInterface

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

 FailureOr<Value> castOrReallocMemRefValue(OpBuilder &b, Value value,
                                           MemRefType type);
 
 /// Try to fold to_memref(to_tensor(x)). If x's type and the result type of the
 /// to_memref op are different, a memref.cast is needed.
 LogicalResult foldToMemrefToTensorPair(RewriterBase &rewriter,
                                        ToMemrefOp toMemref);
 
+/// Return the memory space of the given OpResult as indicated by the optional
+/// `bufferization.memory_space` attribute. Return `None` if no memory space is
+/// specified.
+Optional<unsigned> getMemorySpaceAttr(OpResult result);
+
+/// Set the memory space of the given OpResult as indicated via the
+/// `bufferization.memory_space` attribute.
+void setMemorySpaceAttr(OpResult result, Optional<unsigned> value);
+
 } // namespace bufferization
 } // namespace mlir
 
 #endif // MLIR_DIALECT_BUFFERIZATION_IR_BUFFERIZATION_H_

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

let extraClassDeclaration = [{
     /// 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";
+
+    /// Attribute name used to indicate the memory space of buffer allocations.
+    ///
+    /// Note: Only ops that bufferize to an allocation may have this attribute.
+    constexpr const static ::llvm::StringLiteral
+        kMemorySpaceAttrName = "bufferization.memory_space";
+
+    /// Attribute value used to indicate that no memory space is specified for
+    /// the corresponding OpResult.
+    constexpr const static int64_t
+        kNoMemorySpaceSpecifiedAttrValue = -1;
   }];
   let hasOperationAttrVerify = 1;
 }
 
 #endif // BUFFERIZATION_BASE

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

let description = [{
     `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.
 
-    The optional `memory_space` attribute specifies the memory space when
-    bufferizing this op. The memory space is inferred from `copy` if specified.
-    If neither `copy` nor `memory_space` is specified, the default memory space
-    is used during bufferization.
+    The optional `bufferization.memory_space` attribute specifies the memory
+    space when bufferizing this op. The memory space is inferred from `copy` if
+    there is no `memory_space` attribute. If neither `copy` nor `memory_space`
+    is specified, the default memory space is used during bufferization.
 
     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. Sparse tensor allocations should always be used in a local
     construction operation and never escape the function boundary directly.
 
     Example:
 
     ```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>
     return %0 : tensor<?x?xf32, #SparseMatrix>
     ```
   }];
 
   let arguments = (ins Variadic<Index>:$dynamic_sizes,
-                       Optional<AnyTensor>:$copy,
-                       OptionalAttr<UI64Attr>:$memory_space);
+                       Optional<AnyTensor>:$copy);
 
   let results = (outs AnyTensor:$result);
 
   let extraClassDeclaration = [{
     LogicalResult bufferize(RewriterBase &rewriter,
                             const BufferizationOptions &options);
 
     bool isMemoryWrite(OpResult opResult, 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` or `memory_space`.
+    // Build an op without `copy`.
     OpBuilder<(ins "RankedTensorType":$type, "ValueRange":$dynamicSizes)>,
-
-    // Build an op without `memory_space`.
-    OpBuilder<(ins "RankedTensorType":$type, "ValueRange":$dynamicSizes,
-                   "Value":$copy)>,
   ];
 
   let hasCanonicalizer = 1;
   let hasCustomAssemblyFormat = 1;
   let hasVerifier = 1;
 }
 
 //===----------------------------------------------------------------------===//

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

allocTensorOp->setAttr(BufferizationDialect::kEscapeAttrName,
                          b.getBoolArrayAttr({escape}));
 
   // Add 'memory_space' attribute. Not needed if 'copy' operand is specified.
   if (copy)
     return allocTensorOp.getResult();
   FailureOr<BaseMemRefType> copyBufferType = getBufferType(tensor, options);
   if (failed(copyBufferType))
     return failure();
-  allocTensorOp.setMemorySpaceAttr(
-      b.getIntegerAttr(b.getIntegerType(64, /*isSigned=*/false),
-                       copyBufferType->getMemorySpaceAsInt()));
+  setMemorySpaceAttr(allocTensorOp->getOpResult(0),
+                     copyBufferType->getMemorySpaceAsInt());
   return allocTensorOp.getResult();
 }
 
 LogicalResult BufferizableOpInterface::resolveTensorOpOperandConflicts(
     RewriterBase &rewriter, const AnalysisState &state) {
   OpBuilder::InsertionGuard g(rewriter);
   Operation *op = getOperation();
   SmallVector<OpOperand *> outOfPlaceOpOperands;

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

 /// 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;
 
+/// Attribute name used to indicate the memory space of buffer allocations.
+constexpr const ::llvm::StringLiteral
+    BufferizationDialect::kMemorySpaceAttrName;
+
+/// Attribute value used to indicate that no memory space is specified for the
+/// corresponding OpResult.
+constexpr const int64_t BufferizationDialect::kNoMemorySpaceSpecifiedAttrValue;
+
 //===----------------------------------------------------------------------===//
 // Bufferization Dialect Interfaces
 //===----------------------------------------------------------------------===//
 
 namespace {
 struct BufferizationInlinerInterface : public DialectInlinerInterface {
   using DialectInlinerInterface::DialectInlinerInterface;
 
for (const auto &it : llvm::enumerate(arrayAttr)) {
         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();
   }
+  if (attr.getName() == kMemorySpaceAttrName) {
+    auto arrayAttr = attr.getValue().dyn_cast<ArrayAttr>();
+    if (!arrayAttr)
+      return op->emitError() << "'" << kMemorySpaceAttrName
+                             << "' is expected to be an int array attribute";
+    if (arrayAttr.size() != op->getNumResults())
+      return op->emitError()
+             << "'" << kMemorySpaceAttrName
+             << "' has wrong number of elements, expected "
+             << op->getNumResults() << ", got " << arrayAttr.size();
+    auto bufferizableOp = dyn_cast<BufferizableOpInterface>(op);
+    if (!bufferizableOp)
+      return op->emitError() << "'" << kMemorySpaceAttrName
+                             << "' only valid on bufferizable ops";
+    for (const auto &it : llvm::enumerate(arrayAttr)) {
+      auto attr = it.value();
+      auto intAttr = attr.dyn_cast<IntegerAttr>();
+      if (!intAttr)
+        return op->emitError() << "'" << kMemorySpaceAttrName
+                               << "' is expected to be an int array attribute";
+      if (intAttr.getValue() ==
+          BufferizationDialect::kNoMemorySpaceSpecifiedAttrValue)
+        continue;
+      if (!op->getResult(it.index()).getType().isa<TensorType>())
+        return op->emitError() << "'" << kMemorySpaceAttrName
+                               << "' only valid for tensor results";
+      if (!bufferizableOp.bufferizesToAllocation(op->getOpResult(it.index())))
+        return op->emitError() << "'" << kMemorySpaceAttrName
+                               << "' 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

if (shapedType.isDynamicDim(i)) {
       } else {
         assert(shapedType.isa<RankedTensorType>() && "expected tensor");
         dynamicDims.push_back(b.create<tensor::DimOp>(loc, shapedValue, i));
       }
     }
   }
 }
 
+Optional<unsigned> mlir::bufferization::getMemorySpaceAttr(OpResult result) {
+  Operation *op = result.getDefiningOp();
+  if (!op->hasAttr(BufferizationDialect::kMemorySpaceAttrName))
+    return None;
+  auto arrayAttr =
+      op->getAttr(BufferizationDialect::kMemorySpaceAttrName).cast<ArrayAttr>();
+  int64_t memorySpace =
+      arrayAttr[result.getResultNumber()].cast<IntegerAttr>().getInt();
+  if (memorySpace == BufferizationDialect::kNoMemorySpaceSpecifiedAttrValue)
+    return None;
+  return static_cast<unsigned>(memorySpace);
+}
+
+void mlir::bufferization::setMemorySpaceAttr(OpResult result,
+                                             Optional<unsigned> value) {
+  Operation *op = result.getDefiningOp();
+  assert(cast<BufferizableOpInterface>(op).bufferizesToAllocation(result) &&
+         "expected that OpResult bufferizes to an allocation");
+  SmallVector<int64_t> arrayValues(
+      op->getNumResults(),
+      BufferizationDialect::kNoMemorySpaceSpecifiedAttrValue);
+  if (op->hasAttr(BufferizationDialect::kMemorySpaceAttrName))
+    arrayValues = llvm::to_vector(llvm::map_range(
+        op->getAttr(BufferizationDialect::kMemorySpaceAttrName)
+            .cast<ArrayAttr>(),
+        [](Attribute a) { return a.cast<IntegerAttr>().getInt(); }));
+  arrayValues[result.getResultNumber()] =
+      value.has_value()
+          ? value.value()
+          : BufferizationDialect::kNoMemorySpaceSpecifiedAttrValue;
+  op->setAttr(BufferizationDialect::kMemorySpaceAttrName,
+              OpBuilder(op->getContext()).getI64ArrayAttr(arrayValues));
+}
+
 //===----------------------------------------------------------------------===//
 // AllocTensorOp
 //===----------------------------------------------------------------------===//
 
 LogicalResult AllocTensorOp::bufferize(RewriterBase &rewriter,
                                        const BufferizationOptions &options) {
   OpBuilder::InsertionGuard g(rewriter);
   Location loc = getLoc();

 
 FailureOr<BaseMemRefType> AllocTensorOp::getBufferType(
     Value value, const BufferizationOptions &options,
     const DenseMap<Value, BaseMemRefType> &fixedTypes) {
   assert(value == getResult() && "invalid value");
 
   // Compute memory space of this allocation.
   unsigned memorySpace;
-  if (getMemorySpace().has_value()) {
-    memorySpace = *getMemorySpace();
+  if (auto maybeMemorySpace =
+          getMemorySpaceAttr(getOperation()->getOpResult(0))) {
+    memorySpace = *maybeMemorySpace;
   } else if (getCopy()) {
     auto copyBufferType =
         bufferization::getBufferType(getCopy(), options, fixedTypes);
     if (failed(copyBufferType))
       return failure();
     memorySpace = copyBufferType->getMemorySpaceAsInt();
   } else if (options.defaultMemorySpace.has_value()) {
     memorySpace = *options.defaultMemorySpace;
for (auto &use : getOperation()->getUses())
         return emitError("sparse tensor allocation should not escape function");
   }
 
   return success();
 }
 
 void AllocTensorOp::build(OpBuilder &builder, OperationState &result,
                           RankedTensorType type, ValueRange dynamicSizes) {
-  build(builder, result, type, dynamicSizes, /*copy=*/Value(),
-        /*memory_space=*/IntegerAttr());
-}
-
-void AllocTensorOp::build(OpBuilder &builder, OperationState &result,
-                          RankedTensorType type, ValueRange dynamicSizes,
-                          Value copy) {
-  build(builder, result, type, dynamicSizes, copy,
-        /*memory_space=*/IntegerAttr());
+  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) {
       }
     }
     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());
+    if (auto maybeMemorySpace = getMemorySpaceAttr(op->getOpResult(0)))
+      setMemorySpaceAttr(newOp->getOpResult(0), *maybeMemorySpace);
     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

for (OpOperand &opOperand : bufferizableOp->getOpOperands()) {
       if (usesInsideRegion.empty())
         continue;
 
       // Insert a tensor copy and replace all uses inside of repetitive regions.
       rewriter.setInsertionPoint(bufferizableOp);
       auto tensorCopy = rewriter.create<AllocTensorOp>(
           bufferizableOp->getLoc(), operand.getType().cast<TensorType>(),
           /*dynamicSizes=*/ValueRange(),
-          /*copy=*/operand, /*memory_space=*/IntegerAttr());
+          /*copy=*/operand);
       for (OpOperand *use : usesInsideRegion)
         use->set(tensorCopy);
     }
 
     return WalkResult::advance();
   });
 }
 

mlir/test/Dialect/Arith/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() {memory_space = 0 : ui64} : tensor<10xf32>
-  %1 = bufferization.alloc_tensor() {memory_space = 1 : ui64} : tensor<10xf32>
+  %0 = bufferization.alloc_tensor() {bufferization.memory_space = [0]} : tensor<10xf32>
+  %1 = bufferization.alloc_tensor() {bufferization.memory_space = [1]} : tensor<10xf32>
   // expected-error @+2 {{inconsistent memory space on true/false operands}}
   // expected-error @+1 {{failed to bufferize op}}
   %r = arith.select %c, %0, %1 : tensor<10xf32>
   func.return %r : tensor<10xf32>
 }
 
 // -----
 
 func.func @constant_memory_space(%idx: index, %v: i32) -> tensor<3xi32> {
   // expected-error @+2 {{memory space not implemented yet}}
   // expected-error @+1 {{failed to bufferize op}}
   %cst = arith.constant dense<[5, 1000, 20]> : tensor<3xi32>
   %0 = tensor.insert %v into %cst[%idx] : tensor<3xi32>
   return %0 : tensor<3xi32>
 }
 No newline at end of file

mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize.mlir

func.func @alloc_tensor_with_copy(%t: tensor<5xf32>) -> tensor<5xf32> {
   return %0 : tensor<5xf32>
 }
 
 // -----
 
 // CHECK-LABEL: func @alloc_tensor_with_memory_space()
 func.func @alloc_tensor_with_memory_space() -> tensor<5xf32> {
   // CHECK: %[[alloc:.*]] = memref.alloc() {{.*}} : memref<5xf32, 1>
-  %0 = bufferization.alloc_tensor() {memory_space = 1 : ui64} : tensor<5xf32>
+  %0 = bufferization.alloc_tensor() {bufferization.memory_space = [1]} : tensor<5xf32>
   // CHECK: %[[r:.*]] = bufferization.to_tensor %[[alloc]]
   // CHECK: memref.dealloc %[[alloc]]
   // CHECK: return %[[r]]
   return %0 : tensor<5xf32>
 }
 
 // -----
 

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

 }
 
 // -----
 
 // 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], memory_space = 1 : ui64} : tensor<10xf32>
-  %t = bufferization.alloc_tensor() {memory_space = 1 : ui64} : tensor<10xf32>
-  // CHECK: bufferization.alloc_tensor() {bufferization.escape = [false], memory_space = 1 : ui64} : tensor<10xf32>
+  // 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], memory_space = 0 : ui64} : 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], memory_space = 0 : ui64} : 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], memory_space = 0 : ui64} : 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], memory_space = 0 : ui64} : 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>
 }

mlir/test/Dialect/Bufferization/invalid.mlir

func.func @sparse_alloc_call() {
   // expected-error @+1{{sparse tensor allocation should not escape function}}
   %0 = bufferization.alloc_tensor() : tensor<20x40xf32, #DCSR>
   call @foo(%0) : (tensor<20x40xf32, #DCSR>) -> ()
   return
 }
 
 // -----
 
-func.func @alloc_tensor_invalid_memory_space_attr(%sz: index) {
-  // expected-error @+1{{'bufferization.alloc_tensor' op attribute 'memory_space' failed to satisfy constraint: 64-bit unsigned integer attribute}}
-  %0 = bufferization.alloc_tensor(%sz) {memory_space = "foo"} : tensor<?xf32>
+func.func @invalid_memory_space_attr_type(%sz: index) {
+  // expected-error @+1{{'bufferization.memory_space' is expected to be an int array attribute}}
+  %0 = bufferization.alloc_tensor(%sz) {bufferization.memory_space = "foo"} : tensor<?xf32>
   return
 }
 
+// -----
+
+func.func @invalid_memory_space_attr_size(%sz: index) {
+  // expected-error @+1{{'bufferization.memory_space' has wrong number of elements, expected 1, got 2}}
+  %0 = bufferization.alloc_tensor(%sz) {bufferization.memory_space = [1, 2]} : tensor<?xf32>
+  return
+}
+
+// -----
+
+func.func @invalid_memory_space_attr_op(%t: tensor<?xf32>, %pos: index) {
+  // expected-error @+1{{'bufferization.memory_space' only valid for allocation results}}
+  tensor.extract_slice %t[%pos][5][1] {bufferization.memory_space = [1]} : tensor<?xf32> to tensor<5xf32>
+  return
+}

mlir/test/Dialect/SCF/one-shot-bufferize-invalid.mlir

 // RUN: mlir-opt %s -one-shot-bufferize -split-input-file -verify-diagnostics
 
 func.func @inconsistent_memory_space_scf_if(%c: i1) -> tensor<10xf32> {
   // Yielding tensors with different memory spaces. Such IR cannot be
   // bufferized.
-  %0 = bufferization.alloc_tensor() {memory_space = 0 : ui64} : tensor<10xf32>
-  %1 = bufferization.alloc_tensor() {memory_space = 1 : ui64} : tensor<10xf32>
+  %0 = bufferization.alloc_tensor() {bufferization.memory_space = [0]} : tensor<10xf32>
+  %1 = bufferization.alloc_tensor() {bufferization.memory_space = [1]} : tensor<10xf32>
   // expected-error @+1 {{inconsistent memory space on then/else branches}}
   %r = scf.if %c -> tensor<10xf32> {
     // expected-error @+1 {{failed to bufferize op}}
     scf.yield %0 : tensor<10xf32>
   } else {
     scf.yield %1 : tensor<10xf32>
   }
   func.return %r : tensor<10xf32>
 }

mlir/test/Dialect/SCF/one-shot-bufferize.mlir

 
 // -----
 
 // CHECK-LABEL: func @scf_if_memory_space
 func.func @scf_if_memory_space(%c: i1, %f: f32) -> (f32, f32)
 {
   %c0 = arith.constant 0 : index
   // CHECK: %[[alloc:.*]] = memref.alloc() {{.*}} : memref<5xf32, 1>
-  %0 = bufferization.alloc_tensor() {memory_space = 1 : ui64} : tensor<5xf32>
+  %0 = bufferization.alloc_tensor() {bufferization.memory_space = [1]} : tensor<5xf32>
   // CHECK: scf.if %{{.*}} -> (memref<5xf32, 1>) {
   %1 = scf.if %c -> tensor<5xf32> {
     // CHECK: %[[cloned:.*]] = bufferization.clone %[[alloc]]
     // CHECK: scf.yield %[[cloned]]
     scf.yield %0 : tensor<5xf32>
   } else {
     // CHECK: %[[alloc2:.*]] = memref.alloc() {{.*}} : memref<5xf32, 1>
     // CHECK: memref.store %{{.*}}, %[[alloc2]]

 // CHECK-LABEL: func @scf_execute_region_memory_space
 // CHECK: memref.alloc() {{.*}} : memref<5xf32, 1>
 // CHECK: memref.store
 // CHECK: memref.load
 // CHECK: memref.dealloc
 func.func @scf_execute_region_memory_space(%f: f32) -> f32 {
   %c0 = arith.constant 0 : index
   %0 = scf.execute_region -> tensor<5xf32> {
-    %1 = bufferization.alloc_tensor() {memory_space = 1 : ui64} : tensor<5xf32>
+    %1 = bufferization.alloc_tensor() {bufferization.memory_space = [1]} : tensor<5xf32>
     %2 = tensor.insert %f into %1[%c0] : tensor<5xf32>
     scf.yield %2 : tensor<5xf32>
   }
   %r = tensor.extract %0[%c0] : tensor<5xf32>
   return %r : f32
 }
 
 // -----
 
 // Additional allocs are inserted in the loop body. We just check that all
 // allocs have the correct memory space.
 
 // CHECK-LABEL: func @scf_for_swapping_yields_memory_space
 func.func @scf_for_swapping_yields_memory_space(
     %sz: index, %C : tensor<4xf32>, %lb : index, %ub : index, %step : index)
   -> (f32, f32)
 {
   // CHECK: memref.alloc(%{{.*}}) {{.*}} : memref<?xf32, 1>
   // CHECK: memref.alloc(%{{.*}}) {{.*}} : memref<?xf32, 1>
-  %A = bufferization.alloc_tensor(%sz) {memory_space = 1 : ui64} : tensor<?xf32>
-  %B = bufferization.alloc_tensor(%sz) {memory_space = 1 : ui64} : tensor<?xf32>
+  %A = bufferization.alloc_tensor(%sz) {bufferization.memory_space = [1]} : tensor<?xf32>
+  %B = bufferization.alloc_tensor(%sz) {bufferization.memory_space = [1]} : tensor<?xf32>
 
   // CHECK: scf.for {{.*}} {
   %r0:2 = scf.for %i = %lb to %ub step %step iter_args(%tA = %A, %tB = %B)
       -> (tensor<?xf32>, tensor<?xf32>)
   {
     // CHECK: memref.alloc(%{{.*}}) {{.*}} : memref<?xf32, 1>
     // CHECK: memref.alloc(%{{.*}}) {{.*}} : memref<?xf32, 1>
     %ttA = tensor.insert_slice %C into %tA[0][4][1] : tensor<4xf32> into tensor<?xf32>

mlir/test/Dialect/SparseTensor/sparse_sddmm.mlir

indexing_maps = [
     affine_map<(d0, d1) -> (d0, d1)>,
     affine_map<(d0, d1) -> (d0, d1)>
   ],
   iterator_types = ["parallel", "parallel"]
 }
 
 // CHECK-LABEL: func.func @fold_yield_arg_zero() -> tensor<1024x1024xf64> {
 // CHECK:         %[[VAL_0:.*]] = arith.constant dense<0.000000e+00> : tensor<1024x1024xf64>
-// CHECK:         %[[VAL_1:.*]] = bufferization.alloc_tensor() copy(%[[VAL_0]]) {bufferization.escape = [false], memory_space = 0 : ui64} : tensor<1024x1024xf64>
+// CHECK:         %[[VAL_1:.*]] = bufferization.alloc_tensor() copy(%[[VAL_0]]) {bufferization.escape = [false], bufferization.memory_space = [0]} : tensor<1024x1024xf64>
 // CHECK:         return %[[VAL_1]] : tensor<1024x1024xf64>
 // CHECK:       }
 func.func @fold_yield_arg_zero() -> tensor<1024x1024xf64> {
   %cst = arith.constant 0.000000e+00 : f64
   %0 = tensor.empty() : tensor<1024x1024xf64>
   %1 = linalg.generic {indexing_maps = [affine_map<(d0, d1) -> ()>,
                                         affine_map<(d0, d1) -> (d0, d1)>],
                                         iterator_types = ["parallel", "parallel"]}
                                         ins(%cst : f64)
                                         outs(%0 : tensor<1024x1024xf64>) {
     ^bb0(%a: f64, %x: f64):
       linalg.yield %a : f64
     } -> tensor<1024x1024xf64>
   return %1 : tensor<1024x1024xf64>
 }
 
 // CHECK-LABEL: func.func @fold_yield_direct_zero() -> tensor<32xf64> {
 // CHECK:         %[[VAL_0:.*]] = arith.constant dense<0.000000e+00> : tensor<32xf64>
-// CHECK:         %[[VAL_1:.*]] = bufferization.alloc_tensor() copy(%[[VAL_0]]) {bufferization.escape = [false], memory_space = 0 : ui64} : tensor<32xf64>
+// CHECK:         %[[VAL_1:.*]] = bufferization.alloc_tensor() copy(%[[VAL_0]]) {bufferization.escape = [false], bufferization.memory_space = [0]} : tensor<32xf64>
 // CHECK:         return %[[VAL_1]] : tensor<32xf64>
 // CHECK:       }
 func.func @fold_yield_direct_zero() -> tensor<32xf64> {
   %cst = arith.constant 0.000000e+00 : f64
   %0 = tensor.empty() : tensor<32xf64>
   %1 = linalg.generic {indexing_maps = [affine_map<(d0) -> (d0)>],
                                         iterator_types = ["parallel"]}
                                         outs(%0 : tensor<32xf64>) {
     ^bb0(%x: f64):
       linalg.yield %cst : f64
     } -> tensor<32xf64>
   return %1 : tensor<32xf64>
 }
 
 // CHECK-LABEL: func.func @sampled_dd_unfused(
 // CHECK-SAME:    %[[VAL_0:.*]]: tensor<8x8xf64, #sparse_tensor.encoding<{{.*}}>>,
 // CHECK-SAME:    %[[VAL_1:.*]]: tensor<8x8xf64>,
 // CHECK-SAME:    %[[VAL_2:.*]]: tensor<8x8xf64>) -> tensor<8x8xf64> {
 // CHECK-DAG:     %[[VAL_3:.*]] = arith.constant 8 : index
 // CHECK-DAG:     %[[VAL_4:.*]] = arith.constant 0 : index
 // CHECK-DAG:     %[[VAL_5:.*]] = arith.constant 1 : index
 // CHECK-DAG:     %[[VAL_6:.*]] = arith.constant dense<0.000000e+00> : tensor<8x8xf64>
 // CHECK:         %[[VAL_7:.*]] = bufferization.alloc_tensor() copy(%[[VAL_6]]) {bufferization.escape = [false]} : tensor<8x8xf64>
-// CHECK:         %[[VAL_8:.*]] = bufferization.alloc_tensor() copy(%[[VAL_6]]) {bufferization.escape = [false], memory_space = 0 : ui64} : tensor<8x8xf64>
+// CHECK:         %[[VAL_8:.*]] = bufferization.alloc_tensor() copy(%[[VAL_6]]) {bufferization.escape = [false], bufferization.memory_space = [0]} : tensor<8x8xf64>
 // CHECK:         %[[VAL_9:.*]] = bufferization.to_memref %[[VAL_1]] : memref<8x8xf64>
 // CHECK:         %[[VAL_10:.*]] = bufferization.to_memref %[[VAL_2]] : memref<8x8xf64>
 // CHECK:         %[[VAL_11:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 0 : index} : tensor<8x8xf64, #sparse_tensor.encoding<{{.*}}>> to memref<?xindex>
 // CHECK:         %[[VAL_12:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 0 : index} : tensor<8x8xf64, #sparse_tensor.encoding<{{.*}}>> to memref<?xindex>
 // CHECK:         %[[VAL_13:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 1 : index} : tensor<8x8xf64, #sparse_tensor.encoding<{{.*}}>> to memref<?xindex>
 // CHECK:         %[[VAL_14:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 1 : index} : tensor<8x8xf64, #sparse_tensor.encoding<{{.*}}>> to memref<?xindex>
 // CHECK:         %[[VAL_15:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<8x8xf64, #sparse_tensor.encoding<{{.*}}>> to memref<?xf64>
 // CHECK:         %[[VAL_16:.*]] = bufferization.to_memref %[[VAL_8]] : memref<8x8xf64>