[Mlir-commits] [mlir] [MLIR][Bufferization] Choose default memory space in tensor copy insertion (PR #88500)

[Kunwar Grover]

https://github.com/Groverkss created https://github.com/llvm/llvm-project/pull/88500

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.

>From f4c066734107a334fd06e7e3174e092e0d234770 Mon Sep 17 00:00:00 2001
From: Kunwar Grover <groverkss at gmail.com>
Date: Wed, 10 Apr 2024 16:52:07 +0200
Subject: [PATCH 1/2] fix bufferization default memory space creation in alloc
 tensor

---
 .../Bufferization/IR/BufferizableOpInterface.cpp      | 11 +++++++----
 1 file changed, 7 insertions(+), 4 deletions(-)

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();
 }
 

>From d5f674ad6c679ddffba19276748e90f55f750735 Mon Sep 17 00:00:00 2001
From: Kunwar Grover <groverkss at gmail.com>
Date: Fri, 12 Apr 2024 13:03:24 +0200
Subject: [PATCH 2/2] add test

---
 ...tensor-copy-insertion-memory-space-default.mlir | 14 ++++++++++++++
 1 file changed, 14 insertions(+)
 create mode 100644 mlir/test/Dialect/Bufferization/Transforms/tensor-copy-insertion-memory-space-default.mlir

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>
+}