[MLIR][Bufferization] Choose default memory space in tensor copy insertion

[Groverkss]

Tensor copy insertion currently uses memory_space = 0 when creating a tensor copy using alloc_tensor. This memory space should instead be the default memory space provided in bufferization options.

[llvmbot]

@llvm/pr-subscribers-mlir-sparse
@llvm/pr-subscribers-mlir-bufferization

@llvm/pr-subscribers-mlir

Author: Kunwar Grover (Groverkss)

Changes

Tensor copy insertion currently uses memory_space = 0 when creating a tensor copy using alloc_tensor. This memory space should instead be the default memory space provided in bufferization options.

---

Full diff: https://github.com/llvm/llvm-project/pull/88500.diff

2 Files Affected:

• (modified) mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp (+7-4)
• (added) mlir/test/Dialect/Bufferization/Transforms/tensor-copy-insertion-memory-space-default.mlir (+14)

diff --git a/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp b/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp
index 55c9299c58effd..46f2639c21cad2 100644
--- a/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp
+++ b/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp
@@ -193,10 +193,13 @@ FailureOr<Value> bufferization::allocateTensorForShapedValue(
   FailureOr<BaseMemRefType> copyBufferType = getBufferType(tensor, options);
   if (failed(copyBufferType))
     return failure();
-  Attribute memorySpace = copyBufferType->getMemorySpace();
-  if (!memorySpace)
-    memorySpace = b.getI64IntegerAttr(0);
-  allocTensorOp.setMemorySpaceAttr(memorySpace);
+  std::optional<Attribute> memorySpace = copyBufferType->getMemorySpace();
+  if (!memorySpace) {
+    memorySpace = options.defaultMemorySpaceFn(tensorType);
+  }
+  if (memorySpace.has_value()) {
+    allocTensorOp.setMemorySpaceAttr(memorySpace.value());
+  }
   return allocTensorOp.getResult();
 }
 
diff --git a/mlir/test/Dialect/Bufferization/Transforms/tensor-copy-insertion-memory-space-default.mlir b/mlir/test/Dialect/Bufferization/Transforms/tensor-copy-insertion-memory-space-default.mlir
new file mode 100644
index 00000000000000..e33c95c3710f85
--- /dev/null
+++ b/mlir/test/Dialect/Bufferization/Transforms/tensor-copy-insertion-memory-space-default.mlir
@@ -0,0 +1,14 @@
+// RUN: mlir-opt %s -test-tensor-copy-insertion -split-input-file | FileCheck %s
+
+// -----
+
+// CHECK-LABEL: func @alloc_tensor_default_memory_space
+func.func @alloc_tensor_default_memory_space() -> (tensor<10xf32>, tensor<10xf32>) {
+  %c0 = arith.constant 0 : index
+  %cst = arith.constant 0.0 : f32
+  // CHECK: bufferization.alloc_tensor() : tensor<10xf32>
+  %t = bufferization.alloc_tensor() : tensor<10xf32>
+  // CHECK: bufferization.alloc_tensor() : tensor<10xf32>
+  %s = tensor.insert %cst into %t[%c0] : tensor<10xf32>
+  return %s, %t : tensor<10xf32>, tensor<10xf32>
+}