[Mlir-commits] [mlir] [mlir][tensor][bufferization] Preserve memory space for tensor.concat (PR #213528)

[Hyunjun Shin]

https://github.com/Bammuri updated https://github.com/llvm/llvm-project/pull/213528

>From 965fe9f7c345309fea4407bd5138f902034ae1aa Mon Sep 17 00:00:00 2001
From: Hyunjun Shin <shjj1504 at gmail.com>
Date: Sun, 2 Aug 2026 19:59:26 +0900
Subject: [PATCH] [mlir][tensor][bufferization] Preserve memory space for
 tensor.concat

ConcatOpInterface allocates its destination through allocateTensorForShapedValue, which may record a non-default memory space on the resulting alloc_tensor. It then reconstructs the destination MemRefType without that memory space.

This can make ToBufferOp observe a non-default source buffer and a default-space requested result type. Because the types cannot be cast across memory spaces, bufferization may introduce an additional default-space allocation and a full-result copy.

Query the authoritative buffer type of the alloc_tensor and pass that complete type to ToBufferOp. This also removes the defaultMemorySpaceFn guard, which does not necessarily reflect buffer types selected by other conversion hooks.

Add an encoding-based regression test that checks the destination allocation and both input-to-subview copies remain in the non-default memory space.
---
 .../BufferizableOpInterfaceImpl.cpp           | 18 ++++++----------
 .../Tensor/one-shot-bufferize-encodings.mlir  | 21 +++++++++++++++++++
 2 files changed, 27 insertions(+), 12 deletions(-)

diff --git a/mlir/lib/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.cpp b/mlir/lib/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.cpp
index 525d57341cee5..930f0663e2ab8 100644
--- a/mlir/lib/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.cpp
@@ -1119,18 +1119,12 @@ struct ConcatOpInterface
     if (failed(tensorAlloc))
       return failure();
     auto tensorType = cast<RankedTensorType>(tensorAlloc->getType());
-
-    // TODO: Implement memory space for this op.
-    if (options.defaultMemorySpaceFn(cast<TensorLikeType>(tensorType)) !=
-        Attribute())
-      return op->emitError("memory space not implemented yet");
-
-    MemRefLayoutAttrInterface layout;
-    MemRefType memrefType =
-        MemRefType::get(concatOp.getResultType().getShape(),
-                        concatOp.getResultType().getElementType(), layout);
+    FailureOr<BufferLikeType> memrefType =
+        bufferization::getBufferType(*tensorAlloc, options, state);
+    if (failed(memrefType))
+      return failure();
     Value dstBuffer = bufferization::ToBufferOp::create(
-        rewriter, op->getLoc(), memrefType, *tensorAlloc);
+        rewriter, op->getLoc(), *memrefType, *tensorAlloc);
 
     // Extract the dimension for the concat op
     uint64_t concatDim = concatOp.getDim();
@@ -1166,7 +1160,7 @@ struct ConcatOpInterface
       sizes[concatDim] = concatDimSize;
 
       // Create a subview of the destination buffer.
-      auto dstMemrefType = cast<MemRefType>(memrefType);
+      auto dstMemrefType = cast<MemRefType>(*memrefType);
       MemRefType subviewMemRefType =
           memref::SubViewOp::inferRankReducedResultType(
               operandTensorType.getShape(), dstMemrefType, offsets, sizes,
diff --git a/mlir/test/Dialect/Tensor/one-shot-bufferize-encodings.mlir b/mlir/test/Dialect/Tensor/one-shot-bufferize-encodings.mlir
index 7398fdf614e1a..5dc301d7f2310 100644
--- a/mlir/test/Dialect/Tensor/one-shot-bufferize-encodings.mlir
+++ b/mlir/test/Dialect/Tensor/one-shot-bufferize-encodings.mlir
@@ -18,3 +18,24 @@ func.func @from_elements(%fill: f32, %f: f32, %idx: index) -> tensor<3xf32, 1> {
 //       CHECK:     memref.store %[[arg1]], %[[alloc]][%[[arg2]]] : memref<3xf32, 1>
 //       CHECK:     %[[v0:.+]] = bufferization.to_tensor %[[alloc]] : memref<3xf32, 1> to tensor<3xf32, 1 : i64>
 //       CHECK:     return %[[v0]] : tensor<3xf32, 1 : i64>
+
+// -----
+
+func.func @concat_memory_space(%arg0: tensor<8xf32, 1>,
+                               %arg1: tensor<8xf32, 1>) -> tensor<16xf32, 1> {
+  %0 = tensor.concat dim(0) %arg0, %arg1
+      : (tensor<8xf32, 1>, tensor<8xf32, 1>) -> tensor<16xf32, 1>
+  return %0 : tensor<16xf32, 1>
+}
+
+// CHECK-LABEL: @concat_memory_space
+//  CHECK-SAME: (%[[ARG0:.+]]: tensor<8xf32, 1 : i64>, %[[ARG1:.+]]: tensor<8xf32, 1 : i64>)
+//   CHECK-DAG: %[[BUFFER0:.+]] = bufferization.to_buffer %[[ARG0]]
+//   CHECK-DAG: %[[BUFFER1:.+]] = bufferization.to_buffer %[[ARG1]]
+//       CHECK: %[[ALLOC:.+]] = memref.alloc() {{.*}} : memref<16xf32, 1>
+//       CHECK: %[[SUBVIEW0:.+]] = memref.subview %[[ALLOC]][0] [8] [1]
+//       CHECK: memref.copy %[[BUFFER0]], %[[SUBVIEW0]]
+//       CHECK: %[[SUBVIEW1:.+]] = memref.subview %[[ALLOC]][8] [8] [1]
+//       CHECK: memref.copy %[[BUFFER1]], %[[SUBVIEW1]]
+//       CHECK: %[[RESULT:.+]] = bufferization.to_tensor %[[ALLOC]]
+//       CHECK: return %[[RESULT]] : tensor<16xf32, 1 : i64>