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

[llvmlistbot at llvm.org]

Author: Hyunjun Shin
Date: 2026-08-23T08:37:04Z
New Revision: 62b1dce2e660c48407289bc921871590e68e83a9

URL: https://github.com/llvm/llvm-project/commit/62b1dce2e660c48407289bc921871590e68e83a9
DIFF: https://github.com/llvm/llvm-project/commit/62b1dce2e660c48407289bc921871590e68e83a9.diff

LOG: [mlir][tensor][bufferization] Preserve memory space for tensor.concat (#213528)

## Summary

- Use the buffer type selected for the destination `alloc_tensor` when
bufferizing `tensor.concat`.
- Pass that type to `ToBufferOp` instead of reconstructing a
default-memory-space `MemRefType`.
- Use the same type for the destination subviews, preserving its memory
space and layout.
- Add an encoding-based regression test for a non-default destination
memory space.

## Motivation

While investigating `tensor.concat` bufferization with a non-default
memory space, I reproduced a mismatch between the destination allocation
type and the type requested by `ToBufferOp`.

The destination was allocated as `memref<16xf32, 2>`, but
`ConcatOpInterface` reconstructed a default-space `memref<16xf32>`.
Since `memref.cast` cannot change memory spaces, bufferization inserted
an additional default-space allocation and a `memref.copy`. This extra
allocation and copy defeat the intended memory-space placement.

Before this change, `ConcatOpInterface` explicitly rejected non-default
memory spaces and reconstructed the destination type from only the
shape, element type, and layout. That was sufficient for default-space
cases, but it did not account for other type-conversion hooks that can
select a non-default buffer type.

This change obtains the buffer type selected for the destination
`alloc_tensor` with `getBufferType(*tensorAlloc, options, state)` and
passes that type directly to `ToBufferOp`, preserving the selected
memory space and layout.

## Background

`ConcatOpInterface` allocates its destination through
`allocateTensorForShapedValue`. The resulting `alloc_tensor` may
bufferize to a memref in a non-default memory space. The implementation
then manually reconstructed the destination `MemRefType` without a
memory-space attribute.

This could make `ToBufferOp` observe a non-default-space source buffer
and a default-space requested result type. Since those types cannot be
reconciled with `memref.cast`, bufferization could introduce an
additional default-space allocation and a full-result copy. The
input-to-subview copies that implement the concat itself are still
required.

The existing `defaultMemorySpaceFn` guard is not sufficient because the
actual buffer type may be selected through another type-conversion hook.
Querying `getBufferType(*tensorAlloc, options, state)` obtains the
buffer type selected for the allocation and preserves the selected
memory space and other type information such as layout. Passing this
type to `ToBufferOp` and using it for the destination subviews keeps the
allocation, conversion, and copies consistent.

Assisted-by: Codex

Added: 
    

Modified: 
    mlir/lib/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.cpp
    mlir/test/Dialect/Tensor/one-shot-bufferize-encodings.mlir

Removed: 
    


################################################################################
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>