1// RUN: mlir-opt %s -one-shot-bufferize="bufferize-function-boundaries allow-return-allocs promote-buffer-results-to-out-params function-boundary-type-conversion=fully-dynamic-layout-map" -split-input-file | FileCheck %s
2// RUN: mlir-opt %s -one-shot-bufferize="bufferize-function-boundaries allow-return-allocs promote-buffer-results-to-out-params function-boundary-type-conversion=identity-layout-map" -split-input-file | FileCheck %s --check-prefix=CHECK-NO-LAYOUT
3// RUN: mlir-opt %s -one-shot-bufferize="bufferize-function-boundaries allow-return-allocs function-boundary-type-conversion=infer-layout-map" -split-input-file | FileCheck %s --check-prefix=CHECK-BASELINE
4
5// Note: function-boundary-type-conversion=infer-layout-map with
6// promote-buffer-results-to-out-params is an unsupported combination.
7
8// Note: This bufferization is not very efficient yet, but it works.
9
10// CHECK: #[[$map1:.*]] = affine_map<(d0)[s0, s1] -> (d0 * s1 + s0)>
11// CHECK-LABEL: func @callee(
12//  CHECK-SAME:              %[[arg0:.*]]: memref<5xf32, #[[$map1]]>,
13//  CHECK-SAME:              %[[arg1:.*]]: memref<5xf32, #[[$map1]]>) {
14// This alloc is not needed, but it is inserted due to the out-of-place
15// bufferization of the tensor.insert. With a better layering of the out param
16// promotion pass, this alloc could be avoided.
17//       CHECK:   %[[alloc:.*]] = memref.alloc() {{.*}} : memref<5xf32>
18//       CHECK:   memref.copy %[[arg0]], %[[alloc]]
19//       CHECK:   memref.store %{{.*}}, %[[alloc]]
20//       CHECK:   memref.copy %[[alloc]], %[[arg1]]
21//       CHECK:   memref.dealloc %[[alloc]]
22//       CHECK:   return
23//       CHECK: }
24
25// CHECK-NO-LAYOUT-LABEL: func @callee(
26//  CHECK-NO-LAYOUT-SAME:     %[[arg0:.*]]: memref<5xf32>,
27//  CHECK-NO-LAYOUT-SAME:     %[[arg1:.*]]: memref<5xf32>) {
28//       CHECK-NO-LAYOUT:   %[[alloc:.*]] = memref.alloc() {{.*}} : memref<5xf32>
29//       CHECK-NO-LAYOUT:   memref.copy %[[arg0]], %[[alloc]]
30//       CHECK-NO-LAYOUT:   memref.store {{.*}}, %[[alloc]]
31//       CHECK-NO-LAYOUT:   memref.copy %[[alloc]], %[[arg1]]
32//       CHECK-NO-LAYOUT:   memref.dealloc %[[alloc]]
33
34// CHECK-BASELINE: #[[$map1:.*]] = affine_map<(d0)[s0, s1] -> (d0 * s1 + s0)>
35// CHECK-BASELINE-LABEL: func @callee(
36//  CHECK-BASELINE-SAME:     %[[arg0:.*]]: memref<5xf32, #[[$map1]]>) -> memref<5xf32> {
37//       CHECK-BASELINE:   %[[alloc:.*]] = memref.alloc() {{.*}} : memref<5xf32>
38//       CHECK-BASELINE:   memref.copy %[[arg0]], %[[alloc]]
39//       CHECK-BASELINE:   memref.store {{.*}}, %[[alloc]]
40//       CHECK-BASELINE:   return %[[alloc]]
41func.func @callee(%t: tensor<5xf32>) -> (tensor<5xf32>, tensor<5xf32>) {
42  %c0 = arith.constant 0 : index
43  %cst = arith.constant 8.0 : f32
44  // This must bufferize out-of-place.
45  %1 = tensor.insert %cst into %t[%c0] : tensor<5xf32>
46  // Instead of returning %1, copy into new out param. %t will disappear
47  // entirely because the buffer is equivalent to a bbArg.
48  return %t, %1 : tensor<5xf32>, tensor<5xf32>
49}
50
51// CHECK: func @main(%[[arg0:.*]]: memref<5xf32, #[[$map1]]>) -> (f32, f32) {
52// CHECK:   %[[alloc:.*]] = memref.alloc() : memref<5xf32>
53// CHECK:   %[[casted:.*]] = memref.cast %[[alloc]] : memref<5xf32> to memref<5xf32, #[[$map1]]>
54// CHECK:   call @callee(%[[arg0]], %[[casted]])
55// CHECK:   %[[l1:.*]] = memref.load %[[arg0]]
56// CHECK:   %[[l2:.*]] = memref.load %[[alloc]]
57// CHECK:   memref.dealloc %[[alloc]]
58// CHECK:   return %[[l1]], %[[l2]]
59// CHECK: }
60
61// CHECK-NO-LAYOUT-LABEL: func @main(%{{.*}}: memref<5xf32>) -> (f32, f32) {
62//       CHECK-NO-LAYOUT:   %[[alloc:.*]] = memref.alloc() : memref<5xf32>
63//       CHECK-NO-LAYOUT:   call @callee(%{{.*}}, %[[alloc]])
64func.func @main(%t: tensor<5xf32>) -> (f32, f32) {
65  %c0 = arith.constant 0 : index
66  %0, %1 = func.call @callee(%t)
67      : (tensor<5xf32>) -> (tensor<5xf32>, tensor<5xf32>)
68  %2 = tensor.extract %0[%c0] : tensor<5xf32>
69  %3 = tensor.extract %1[%c0] : tensor<5xf32>
70  return %2, %3 : f32, f32
71}
72
73// -----
74
75// CHECK: #[[$map2a:.*]] = affine_map<(d0, d1)[s0, s1, s2] -> (d0 * s1 + s0 + d1 * s2)>
76// CHECK: #[[$map2b:.*]] = affine_map<(d0, d1)[s0] -> (d0 * 20 + s0 + d1)>
77// CHECK-LABEL: func @callee(
78//  CHECK-SAME:     %{{.*}}: index,
79//  CHECK-SAME:     %[[r:.*]]: memref<2x5xf32, #[[$map2a]]>) {
80//       CHECK:   %[[alloc:.*]] = memref.alloc() {{.*}} : memref<10x20xf32>
81//       CHECK:   %[[subview:.*]] = memref.subview %[[alloc]]{{.*}} : memref<10x20xf32> to memref<2x5xf32, #[[$map2b]]>
82//       CHECK:   memref.copy %[[subview]], %[[r]]
83//       CHECK:   memref.dealloc %[[alloc]]
84
85// CHECK-NO-LAYOUT-LABEL: func @callee(
86//  CHECK-NO-LAYOUT-SAME:              %{{.*}}: index,
87//  CHECK-NO-LAYOUT-SAME:              %[[r:.*]]: memref<2x5xf32>) {
88//       CHECK-NO-LAYOUT:   %[[alloc:.*]] = memref.alloc() {{.*}} : memref<10x20xf32>
89//       CHECK-NO-LAYOUT:   %[[subview:.*]] = memref.subview %[[alloc]]
90// Note: This alloc is not needed, but it is inserted before the returned buffer
91// is promoted to an out param to reconcile mismatching layout maps on return
92// value and function signature.
93//       CHECK-NO-LAYOUT:   %[[alloc2:.*]] = memref.alloc() : memref<2x5xf32>
94//       CHECK-NO-LAYOUT:   memref.copy %[[subview]], %[[alloc2]]
95//       CHECK-NO-LAYOUT:   memref.dealloc %[[alloc]]
96//       CHECK-NO-LAYOUT:   memref.copy %[[alloc2]], %[[r]]
97//       CHECK-NO-LAYOUT:   memref.dealloc %[[alloc2]]
98
99// CHECK-BASELINE: #[[$map2:.*]] = affine_map<(d0, d1)[s0] -> (d0 * 20 + s0 + d1)>
100// CHECK-BASELINE-LABEL: func @callee(
101//  CHECK-BASELINE-SAME:     %{{.*}}: index) -> memref<2x5xf32, #[[$map2]]> {
102//       CHECK-BASELINE:   %[[alloc:.*]] = memref.alloc() {{.*}} : memref<10x20xf32>
103//       CHECK-BASELINE:   %[[subview:.*]] = memref.subview %[[alloc]]
104//       CHECK-BASELINE:   return %[[subview]]
105func.func @callee(%idx: index) -> tensor<2x5xf32> {
106  %0 = bufferization.alloc_tensor() : tensor<10x20xf32>
107  %1 = tensor.extract_slice %0[%idx, %idx][2, 5][1, 1] : tensor<10x20xf32> to tensor<2x5xf32>
108  return %1 : tensor<2x5xf32>
109}
110
111// CHECK: func @main(
112// CHECK:   %[[alloc:.*]] = memref.alloc() : memref<2x5xf32>
113// CHECK:   %[[casted:.*]] = memref.cast %[[alloc]] : memref<2x5xf32> to memref<2x5xf32, #[[$map2a]]>
114// CHECK:   call @callee(%{{.*}}, %[[casted]])
115// CHECK:   memref.load %[[alloc]]
116// CHECK:   memref.dealloc %[[alloc]]
117
118// CHECK-NO-LAYOUT: func @main(
119// CHECK-NO-LAYOUT:   %[[alloc:.*]] = memref.alloc() : memref<2x5xf32>
120// CHECK-NO-LAYOUT:   call @callee(%{{.*}}, %[[alloc]])
121// CHECK-NO-LAYOUT:   memref.load %[[alloc]]
122// CHECK-NO-LAYOUT:   memref.dealloc
123
124// CHECK-BASELINE: func @main(
125// CHECK-BASELINE:   %[[call:.*]] = call @callee
126// CHECK-BASELINE:   memref.load %[[call]]
127func.func @main(%idx: index) -> f32 {
128  %c0 = arith.constant 0 : index
129  %0 = func.call @callee(%idx) : (index) -> (tensor<2x5xf32>)
130  %1 = tensor.extract %0[%c0, %c0] : tensor<2x5xf32>
131  return %1 : f32
132}
133