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

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: HyunJun (Bammuri)

<details>
<summary>Changes</summary>

## Summary

- query the buffer type of the destination `alloc_tensor` when bufferizing `tensor.concat`
- pass the complete type to `ToBufferOp` instead of reconstructing a default-memory-space `MemRefType`
- add an encoding-based regression test for a non-default destination memory space

## 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)` uses the authoritative type of the allocation and also preserves other type information such as layout.

## Testing

- `ninja -j3 tools/mlir/lib/Dialect/Tensor/Transforms/CMakeFiles/obj.MLIRTensorTransforms.dir/BufferizableOpInterfaceImpl.cpp.o`
- `llvm-lit -v mlir/test/Dialect/Tensor/one-shot-bufferize-encodings.mlir mlir/test/Dialect/Tensor/bufferize.mlir`

Both lit tests pass.

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


2 Files Affected:

- (modified) mlir/lib/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.cpp (+6-12) 
- (modified) mlir/test/Dialect/Tensor/one-shot-bufferize-encodings.mlir (+21) 


``````````diff
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>

``````````

</details>


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