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