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

[Kunwar Grover]

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

>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/4] 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/4] 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>
+}

>From 19c96c425cf959c04e4a2ae5befae077a3f397e4 Mon Sep 17 00:00:00 2001
From: Kunwar Grover <groverkss at gmail.com>
Date: Fri, 12 Apr 2024 13:15:44 +0200
Subject: [PATCH 3/4] style fix

---
 .../Dialect/Bufferization/IR/BufferizableOpInterface.cpp    | 6 ++----
 1 file changed, 2 insertions(+), 4 deletions(-)

diff --git a/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp b/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp
index 46f2639c21cad2..c2b2b99fc0083b 100644
--- a/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp
+++ b/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp
@@ -194,12 +194,10 @@ FailureOr<Value> bufferization::allocateTensorForShapedValue(
   if (failed(copyBufferType))
     return failure();
   std::optional<Attribute> memorySpace = copyBufferType->getMemorySpace();
-  if (!memorySpace) {
+  if (!memorySpace)
     memorySpace = options.defaultMemorySpaceFn(tensorType);
-  }
-  if (memorySpace.has_value()) {
+  if (memorySpace.has_value())
     allocTensorOp.setMemorySpaceAttr(memorySpace.value());
-  }
   return allocTensorOp.getResult();
 }
 

>From 33f2193583627ce6f1966f9e2920c93e1a2e58bf Mon Sep 17 00:00:00 2001
From: Kunwar Grover <groverkss at gmail.com>
Date: Fri, 12 Apr 2024 13:46:23 +0200
Subject: [PATCH 4/4] fix tests

---
 ...tensor-copy-insertion-memory-space-default.mlir | 14 --------------
 .../Transforms/tensor-copy-insertion.mlir          |  6 +++---
 mlir/test/Dialect/SparseTensor/sparse_sddmm.mlir   |  2 +-
 3 files changed, 4 insertions(+), 18 deletions(-)
 delete 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
deleted file mode 100644
index e33c95c3710f85..00000000000000
--- a/mlir/test/Dialect/Bufferization/Transforms/tensor-copy-insertion-memory-space-default.mlir
+++ /dev/null
@@ -1,14 +0,0 @@
-// 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>
-}
diff --git a/mlir/test/Dialect/Bufferization/Transforms/tensor-copy-insertion.mlir b/mlir/test/Dialect/Bufferization/Transforms/tensor-copy-insertion.mlir
index 72cf08df5978cf..55dcca193cbe49 100644
--- a/mlir/test/Dialect/Bufferization/Transforms/tensor-copy-insertion.mlir
+++ b/mlir/test/Dialect/Bufferization/Transforms/tensor-copy-insertion.mlir
@@ -34,7 +34,7 @@ func.func @do_not_copy_undefined_tensor(%f: f32, %idx: index)
 {
   // The second alloc_tensor should not have a copy operand.
   // CHECK: bufferization.alloc_tensor() : tensor<5xf32>
-  // CHECK: bufferization.alloc_tensor() {memory_space = 0 : i64} : tensor<5xf32>
+  // CHECK: bufferization.alloc_tensor() : tensor<5xf32>
   %0 = bufferization.alloc_tensor() : tensor<5xf32>
   %1 = tensor.insert %f into %0[%idx] : tensor<5xf32>
   return %0, %1 : tensor<5xf32>, tensor<5xf32>
@@ -46,7 +46,7 @@ func.func @do_not_copy_undefined_tensor(%f: f32, %idx: index)
 func.func @do_not_copy_when_overwritten(%t: tensor<5xf32>, %f: f32)
   -> (tensor<5xf32>, tensor<5xf32>)
 {
-  // CHECK: %[[alloc:.*]] = bufferization.alloc_tensor() {memory_space = 0 : i64} : tensor<5xf32>
+  // CHECK: %[[alloc:.*]] = bufferization.alloc_tensor() : tensor<5xf32>
   // CHECK: linalg.generic {{.*}} outs(%[[alloc]] : tensor<5xf32>)
   %r = linalg.generic {
     indexing_maps = [affine_map<(d0) -> (d0)>],
@@ -65,7 +65,7 @@ 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() {memory_space = 0 : i64} : tensor<3xf32>
+  // CHECK: %[[alloc:.*]] = bufferization.alloc_tensor() : tensor<3xf32>
   // CHECK: linalg.generic {{.*}} outs(%[[alloc]] : tensor<3xf32>)
   %r = linalg.generic {
     indexing_maps = [affine_map<(d0) -> (d0)>],
diff --git a/mlir/test/Dialect/SparseTensor/sparse_sddmm.mlir b/mlir/test/Dialect/SparseTensor/sparse_sddmm.mlir
index 6fc2db481c6fd4..fcd69bea426d67 100644
--- a/mlir/test/Dialect/SparseTensor/sparse_sddmm.mlir
+++ b/mlir/test/Dialect/SparseTensor/sparse_sddmm.mlir
@@ -63,7 +63,7 @@ func.func @fold_yield_direct_zero() -> tensor<32xf64> {
 // CHECK-DAG:       %[[VAL_5:.*]] = arith.constant 1 : index
 // CHECK-DAG:       %[[VAL_6:.*]] = arith.constant dense<0.000000e+00> : tensor<8x8xf64>
 // CHECK-DAG:       %[[VAL_7:.*]] = bufferization.alloc_tensor() copy(%[[VAL_6]]) : tensor<8x8xf64>
-// CHECK-DAG:       %[[VAL_8:.*]] = bufferization.alloc_tensor() copy(%[[VAL_6]]) {memory_space = 0 : i64} : tensor<8x8xf64>
+// CHECK-DAG:       %[[VAL_8:.*]] = bufferization.alloc_tensor() copy(%[[VAL_6]]) : tensor<8x8xf64>
 // CHECK-DAG:       %[[VAL_9:.*]] = bufferization.to_memref %[[VAL_1]] : memref<8x8xf64>
 // CHECK-DAG:       %[[VAL_10:.*]] = bufferization.to_memref %[[VAL_2]] : memref<8x8xf64>
 // CHECK-DAG:       %[[VAL_11:.*]] = sparse_tensor.positions %[[VAL_0]] {level = 0 : index} : tensor<8x8xf64, #sparse{{[0-9]*}}> to memref<?xindex>