1// RUN: mlir-opt %s -func-bufferize -split-input-file -verify-diagnostics | FileCheck %s
2
3// CHECK-LABEL:   func @identity(
4// CHECK-SAME:                   %[[ARG:.*]]: memref<f32>) -> memref<f32> {
5// CHECK:           return %[[ARG]] : memref<f32>
6func.func @identity(%arg0: tensor<f32>) -> tensor<f32> {
7  return %arg0 : tensor<f32>
8}
9
10// CHECK-LABEL:   func @block_arguments(
11// CHECK-SAME:        %[[ARG:.*]]: memref<f32>) -> memref<f32> {
12// CHECK:           cf.br ^bb1(%[[ARG]] : memref<f32>)
13// CHECK:         ^bb1(%[[BBARG:.*]]: memref<f32>):
14// CHECK:           return %[[BBARG]] : memref<f32>
15func.func @block_arguments(%arg0: tensor<f32>) -> tensor<f32> {
16  cf.br ^bb1(%arg0: tensor<f32>)
17^bb1(%bbarg: tensor<f32>):
18  return %bbarg : tensor<f32>
19}
20
21// CHECK-LABEL:   func private @source() -> memref<f32>
22// CHECK-LABEL:   func @call_source() -> memref<f32> {
23// CHECK:           %[[RET:.*]] = call @source() : () -> memref<f32>
24// CHECK:           return %[[RET]] : memref<f32>
25func.func private @source() -> tensor<f32>
26func.func @call_source() -> tensor<f32> {
27  %0 = call @source() : () -> tensor<f32>
28  return %0 : tensor<f32>
29}
30// CHECK-LABEL:   func @call_sink(
31// CHECK-SAME:                    %[[ARG:.*]]: memref<f32>) {
32// CHECK:           call @sink(%[[ARG]]) : (memref<f32>) -> ()
33// CHECK:           return
34func.func private @sink(tensor<f32>)
35func.func @call_sink(%arg0: tensor<f32>) {
36  call @sink(%arg0) : (tensor<f32>) -> ()
37  return
38}
39
40// CHECK-LABEL:   func @unconverted_op_in_body() -> memref<f32> {
41// CHECK:           %[[TENSOR:.*]] = "test.source"() : () -> tensor<f32>
42// CHECK:           %[[MEMREF:.*]] = bufferization.to_memref %[[TENSOR]] : memref<f32>
43// CHECK:           return %[[MEMREF]] : memref<f32>
44func.func @unconverted_op_in_body() -> tensor<f32> {
45  %0 = "test.source"() : () -> tensor<f32>
46  return %0 : tensor<f32>
47}
48
49// -----
50
51// Because this pass updates block arguments, it needs to also atomically
52// update all terminators and issue an error if that is not possible.
53func.func @unable_to_update_terminator(%arg0: tensor<f32>) -> tensor<f32> {
54    %0 = arith.constant true
55    cf.cond_br %0, ^bb1(%arg0: tensor<f32>), ^bb2(%arg0: tensor<f32>)
56  ^bb1(%bbarg0: tensor<f32>):
57    // expected-error @+1 {{failed to legalize operation 'test.terminator'}}
58    "test.terminator"() : () -> ()
59  ^bb2(%bbarg1: tensor<f32>):
60    return %bbarg1 : tensor<f32>
61}
62
63// -----
64
65// There was a bug in func-bufferize pass which caused terminators without
66// ReturnLike and BranchOpInterface traits (e.g. scf.condition) to always
67// fail to legalize even if bufferization doesn't needed.
68// Check the pass succedeed.
69// CHECK: bufferize_while
70// CHECK: scf.while
71// CHECK: scf.condition
72func.func @bufferize_while(%arg0: i64, %arg1: i64) -> i64 {
73  %c2_i64 = arith.constant 2 : i64
74  %0:2 = scf.while (%arg2 = %arg0) : (i64) -> (i64, i64) {
75    %1 = arith.cmpi slt, %arg2, %arg1 : i64
76    scf.condition(%1) %arg2, %arg2 : i64, i64
77  } do {
78  ^bb0(%arg2: i64, %arg3: i64):
79    %1 = arith.muli %arg3, %c2_i64 : i64
80    scf.yield %1 : i64
81  }
82  return %0#1 : i64
83}
84