xref: /llvm-project/mlir/test/Dialect/Bufferization/Transforms/tensor-copy-insertion-memory-space.mlir (revision 6bf043e7433680c6f4e36393734ef83699b30f14) 
1// RUN: mlir-opt %s -test-tensor-copy-insertion="must-infer-memory-space" -split-input-file | FileCheck %s
2
3// CHECK-LABEL: func @unknown_op_copy
4func.func @unknown_op_copy() -> (tensor<10xf32>, tensor<10xf32>) {
5  %c0 = arith.constant 0 : index
6  %cst = arith.constant 0.0 : f32
7  // CHECK: %[[dummy:.*]] = "test.dummy_op"() : () -> tensor<10xf32>
8  %t = "test.dummy_op"() : () -> tensor<10xf32>
9  // CHECK: %[[copy:.*]] = bufferization.alloc_tensor() copy(%[[dummy]]) : tensor<10xf32>
10  %s = tensor.insert %cst into %t[%c0] : tensor<10xf32>
11  return %s, %t : tensor<10xf32>, tensor<10xf32>
12}
13
14// -----
15
16// CHECK-LABEL: func @alloc_tensor_copy
17func.func @alloc_tensor_copy() -> (tensor<10xf32>, tensor<10xf32>) {
18  %c0 = arith.constant 0 : index
19  %cst = arith.constant 0.0 : f32
20  // CHECK: bufferization.alloc_tensor() {memory_space = 1 : ui64} : tensor<10xf32>
21  %t = bufferization.alloc_tensor() {memory_space = 1 : ui64} : tensor<10xf32>
22  // CHECK: bufferization.alloc_tensor() {memory_space = 1 : ui64} : tensor<10xf32>
23  %s = tensor.insert %cst into %t[%c0] : tensor<10xf32>
24  return %s, %t : tensor<10xf32>, tensor<10xf32>
25}
26