xref: /llvm-project/mlir/test/Transforms/inlining.mlir (revision d072ca1a496cc3f4ad0adf6f7d43f76406a704d6)
1// RUN: mlir-opt %s -inline='default-pipeline=''' | FileCheck %s
2// RUN: mlir-opt %s --mlir-disable-threading -inline='default-pipeline=''' | FileCheck %s
3// RUN: mlir-opt %s -inline='default-pipeline=''' -mlir-print-debuginfo -mlir-print-local-scope | FileCheck %s --check-prefix INLINE-LOC
4// RUN: mlir-opt %s -inline | FileCheck %s --check-prefix INLINE_SIMPLIFY
5// RUN: mlir-opt %s -inline='op-pipelines=func.func(canonicalize,cse)' | FileCheck %s --check-prefix INLINE_SIMPLIFY
6
7// Inline a function that takes an argument.
8func.func @func_with_arg(%c : i32) -> i32 {
9  %b = arith.addi %c, %c : i32
10  return %b : i32
11}
12
13// CHECK-LABEL: func @inline_with_arg
14func.func @inline_with_arg(%arg0 : i32) -> i32 {
15  // CHECK-NEXT: arith.addi
16  // CHECK-NEXT: return
17
18  %0 = call @func_with_arg(%arg0) : (i32) -> i32
19  return %0 : i32
20}
21
22// CHECK-LABEL: func @noinline_with_arg
23func.func @noinline_with_arg(%arg0 : i32) -> i32 {
24  // CHECK-NEXT: func_with_arg
25  // CHECK-NEXT: return
26
27  %0 = call @func_with_arg(%arg0) {no_inline} : (i32) -> i32
28  return %0 : i32
29}
30
31func.func @non_inlinable_func_with_arg(%c : i32) -> i32 attributes {no_inline} {
32  %b = arith.addi %c, %c : i32
33  return %b : i32
34}
35
36// CHECK-LABEL: func @noinline_with_func_arg
37func.func @noinline_with_func_arg(%arg0 : i32) -> i32 {
38  // CHECK-NEXT: non_inlinable_func_with_arg
39  // CHECK-NEXT: return
40
41  %0 = call @non_inlinable_func_with_arg(%arg0) : (i32) -> i32
42  return %0 : i32
43}
44
45// Inline a function that has multiple return operations.
46func.func @func_with_multi_return(%a : i1) -> (i32) {
47  cf.cond_br %a, ^bb1, ^bb2
48
49^bb1:
50  %const_0 = arith.constant 0 : i32
51  return %const_0 : i32
52
53^bb2:
54  %const_55 = arith.constant 55 : i32
55  return %const_55 : i32
56}
57
58// CHECK-LABEL: func @inline_with_multi_return() -> i32
59func.func @inline_with_multi_return() -> i32 {
60// CHECK-NEXT:    [[VAL_7:%.*]] = arith.constant false
61// CHECK-NEXT:    cf.cond_br [[VAL_7]], ^bb1, ^bb2
62// CHECK:       ^bb1:
63// CHECK-NEXT:    [[VAL_8:%.*]] = arith.constant 0 : i32
64// CHECK-NEXT:    cf.br ^bb3([[VAL_8]] : i32)
65// CHECK:       ^bb2:
66// CHECK-NEXT:    [[VAL_9:%.*]] = arith.constant 55 : i32
67// CHECK-NEXT:    cf.br ^bb3([[VAL_9]] : i32)
68// CHECK:       ^bb3([[VAL_10:%.*]]: i32):
69// CHECK-NEXT:    return [[VAL_10]] : i32
70
71  %false = arith.constant false
72  %x = call @func_with_multi_return(%false) : (i1) -> i32
73  return %x : i32
74}
75
76// Check that location information is updated for inlined instructions.
77func.func @func_with_locations(%c : i32) -> i32 {
78  %b = arith.addi %c, %c : i32 loc("mysource.cc":10:8)
79  return %b : i32 loc("mysource.cc":11:2)
80}
81
82// INLINE-LOC-LABEL: func @inline_with_locations
83func.func @inline_with_locations(%arg0 : i32) -> i32 {
84  // INLINE-LOC-NEXT: arith.addi %{{.*}}, %{{.*}} : i32 loc(callsite("mysource.cc":10:8 at "mysource.cc":55:14))
85  // INLINE-LOC-NEXT: return
86
87  %0 = call @func_with_locations(%arg0) : (i32) -> i32 loc("mysource.cc":55:14)
88  return %0 : i32
89}
90
91
92// Check that external function declarations are not inlined.
93func.func private @func_external()
94
95// CHECK-LABEL: func @no_inline_external
96func.func @no_inline_external() {
97  // CHECK-NEXT: call @func_external()
98  call @func_external() : () -> ()
99  return
100}
101
102// Check that multiple levels of calls will be inlined.
103func.func @multilevel_func_a() {
104  return
105}
106func.func @multilevel_func_b() {
107  call @multilevel_func_a() : () -> ()
108  return
109}
110
111// CHECK-LABEL: func @inline_multilevel
112func.func @inline_multilevel() {
113  // CHECK-NOT: call
114  %fn = "test.functional_region_op"() ({
115    call @multilevel_func_b() : () -> ()
116    "test.return"() : () -> ()
117  }) : () -> (() -> ())
118
119  call_indirect %fn() : () -> ()
120  return
121}
122
123// Check that recursive calls are not inlined.
124// CHECK-LABEL: func @no_inline_recursive
125func.func @no_inline_recursive() {
126  // CHECK: test.functional_region_op
127  // CHECK-NOT: test.functional_region_op
128  %fn = "test.functional_region_op"() ({
129    call @no_inline_recursive() : () -> ()
130    "test.return"() : () -> ()
131  }) : () -> (() -> ())
132  return
133}
134
135// Check that we can convert types for inputs and results as necessary.
136func.func @convert_callee_fn(%arg : i32) -> i32 {
137  return %arg : i32
138}
139func.func @convert_callee_fn_multi_arg(%a : i32, %b : i32) -> () {
140  return
141}
142func.func @convert_callee_fn_multi_res() -> (i32, i32) {
143  %res = arith.constant 0 : i32
144  return %res, %res : i32, i32
145}
146
147// CHECK-LABEL: func @inline_convert_call
148func.func @inline_convert_call() -> i16 {
149  // CHECK: %[[INPUT:.*]] = arith.constant
150  %test_input = arith.constant 0 : i16
151
152  // CHECK: %[[CAST_INPUT:.*]] = "test.cast"(%[[INPUT]]) : (i16) -> i32
153  // CHECK: %[[CAST_RESULT:.*]] = "test.cast"(%[[CAST_INPUT]]) : (i32) -> i16
154  // CHECK-NEXT: return %[[CAST_RESULT]]
155  %res = "test.conversion_call_op"(%test_input) { callee=@convert_callee_fn } : (i16) -> (i16)
156  return %res : i16
157}
158
159func.func @convert_callee_fn_multiblock() -> i32 {
160  cf.br ^bb0
161^bb0:
162  %0 = arith.constant 0 : i32
163  return %0 : i32
164}
165
166// CHECK-LABEL: func @inline_convert_result_multiblock
167func.func @inline_convert_result_multiblock() -> i16 {
168// CHECK:   cf.br ^bb1 {inlined_conversion}
169// CHECK: ^bb1:
170// CHECK:   %[[C:.+]] = arith.constant {inlined_conversion} 0 : i32
171// CHECK:   cf.br ^bb2(%[[C]] : i32)
172// CHECK: ^bb2(%[[BBARG:.+]]: i32):
173// CHECK:   %[[CAST_RESULT:.+]] = "test.cast"(%[[BBARG]]) : (i32) -> i16
174// CHECK:   return %[[CAST_RESULT]] : i16
175
176  %res = "test.conversion_call_op"() { callee=@convert_callee_fn_multiblock } : () -> (i16)
177  return %res : i16
178}
179
180// CHECK-LABEL: func @no_inline_convert_call
181func.func @no_inline_convert_call() {
182  // CHECK: "test.conversion_call_op"
183  %test_input_i16 = arith.constant 0 : i16
184  %test_input_i64 = arith.constant 0 : i64
185  "test.conversion_call_op"(%test_input_i16, %test_input_i64) { callee=@convert_callee_fn_multi_arg } : (i16, i64) -> ()
186
187  // CHECK: "test.conversion_call_op"
188  %res_2:2 = "test.conversion_call_op"() { callee=@convert_callee_fn_multi_res } : () -> (i16, i64)
189  return
190}
191
192// Check that we properly simplify when inlining.
193func.func @simplify_return_constant() -> i32 {
194  %res = arith.constant 0 : i32
195  return %res : i32
196}
197
198func.func @simplify_return_reference() -> (() -> i32) {
199  %res = constant @simplify_return_constant : () -> i32
200  return %res : () -> i32
201}
202
203// INLINE_SIMPLIFY-LABEL: func @inline_simplify
204func.func @inline_simplify() -> i32 {
205  // INLINE_SIMPLIFY-NEXT: %[[CST:.*]] = arith.constant 0 : i32
206  // INLINE_SIMPLIFY-NEXT: return %[[CST]]
207  %fn = call @simplify_return_reference() : () -> (() -> i32)
208  %res = call_indirect %fn() : () -> i32
209  return %res : i32
210}
211
212// CHECK-LABEL: func @no_inline_invalid_call
213func.func @no_inline_invalid_call() -> i32 {
214  %res = "test.conversion_call_op"() { callee=@convert_callee_fn_multiblock, noinline } : () -> (i32)
215  return %res : i32
216}
217
218func.func @gpu_alloc() -> memref<1024xf32> {
219  %m = gpu.alloc [] () : memref<1024xf32>
220  return %m : memref<1024xf32>
221}
222
223// CHECK-LABEL: func @inline_gpu_ops
224func.func @inline_gpu_ops() -> memref<1024xf32> {
225  // CHECK-NEXT: gpu.alloc
226  %m = call @gpu_alloc() : () -> memref<1024xf32>
227  return %m : memref<1024xf32>
228}
229
230// Test block arguments location propagation.
231// Use two call-sites to force cloning.
232func.func @func_with_block_args_location(%arg0 : i32) {
233  cf.br ^bb1(%arg0 : i32)
234^bb1(%x : i32 loc("foo")):
235  "test.foo" (%x) : (i32) -> () loc("bar")
236  return
237}
238
239// INLINE-LOC-LABEL: func @func_with_block_args_location_callee1
240// INLINE-LOC: cf.br
241// INLINE-LOC: ^bb{{[0-9]+}}(%{{.*}}: i32 loc(callsite("foo" at "bar"))
242func.func @func_with_block_args_location_callee1(%arg0 : i32) {
243  call @func_with_block_args_location(%arg0) : (i32) -> () loc("bar")
244  return
245}
246
247// CHECK-LABEL: func @func_with_block_args_location_callee2
248func.func @func_with_block_args_location_callee2(%arg0 : i32) {
249  call @func_with_block_args_location(%arg0) : (i32) -> ()
250  return
251}
252
253func.func @func_with_multiple_blocks(%arg0 : i32) {
254  cf.br ^bb1(%arg0 : i32)
255^bb1(%x : i32):
256  "test.foo" (%x) : (i32) -> () loc("bar")
257  return
258}
259
260// CHECK-LABEL: func @func_with_multiple_blocks_callee1
261func.func @func_with_multiple_blocks_callee1(%arg0 : i32) {
262  "test.dummy_op"() ({
263    // Call cannot be inlined because "test.dummy" may not support unstructured
264    // control flow in its body.
265    // CHECK: call @func_with_multiple_blocks
266    call @func_with_multiple_blocks(%arg0) : (i32) -> ()
267    "test.terminator"() : () -> ()
268  }) : () -> ()
269  return
270}
271
272// CHECK-LABEL: func @func_with_multiple_blocks_callee2
273func.func @func_with_multiple_blocks_callee2(%arg0 : i32, %c : i1) {
274  %0 = scf.while (%arg1 = %arg0) : (i32) -> (i32) {
275    // Call cannot be inlined because scf.while does not support unstructured
276    // control flow in its body.
277    // CHECK: call @func_with_multiple_blocks
278    func.call @func_with_multiple_blocks(%arg0) : (i32) -> ()
279    scf.condition(%c) %arg1 : i32
280  } do {
281  ^bb0(%arg1: i32):
282    scf.yield %arg1 : i32
283  }
284  return
285}
286
287// Check that we can handle argument and result attributes.
288test.conversion_func_op @handle_attr_callee_fn_multi_arg(%arg0 : i16, %arg1 : i16 {"test.handle_argument"}) -> (i16 {"test.handle_result"}, i16) {
289  %0 = arith.addi %arg0, %arg1 : i16
290  %1 = arith.subi %arg0, %arg1 : i16
291  "test.return"(%0, %1) : (i16, i16) -> ()
292}
293test.conversion_func_op @handle_attr_callee_fn(%arg0 : i32 {"test.handle_argument"}) -> (i32 {"test.handle_result"}) {
294  "test.return"(%arg0) : (i32) -> ()
295}
296
297// CHECK-LABEL: func @inline_handle_attr_call
298// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]
299// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]
300func.func @inline_handle_attr_call(%arg0 : i16, %arg1 : i16) -> (i16, i16) {
301
302  // CHECK: %[[CHANGE_INPUT:.*]] = "test.type_changer"(%[[ARG1]]) : (i16) -> i16
303  // CHECK: %[[SUM:.*]] = arith.addi %[[ARG0]], %[[CHANGE_INPUT]]
304  // CHECK: %[[DIFF:.*]] = arith.subi %[[ARG0]], %[[CHANGE_INPUT]]
305  // CHECK: %[[CHANGE_RESULT:.*]] = "test.type_changer"(%[[SUM]]) : (i16) -> i16
306  // CHECK-NEXT: return %[[CHANGE_RESULT]], %[[DIFF]]
307  %res0, %res1 = "test.conversion_call_op"(%arg0, %arg1) { callee=@handle_attr_callee_fn_multi_arg } : (i16, i16) -> (i16, i16)
308  return %res0, %res1 : i16, i16
309}
310
311// CHECK-LABEL: func @inline_convert_and_handle_attr_call
312// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]
313func.func @inline_convert_and_handle_attr_call(%arg0 : i16) -> (i16) {
314
315  // CHECK: %[[CAST_INPUT:.*]] = "test.cast"(%[[ARG0]]) : (i16) -> i32
316  // CHECK: %[[CHANGE_INPUT:.*]] = "test.type_changer"(%[[CAST_INPUT]]) : (i32) -> i32
317  // CHECK: %[[CHANGE_RESULT:.*]] = "test.type_changer"(%[[CHANGE_INPUT]]) : (i32) -> i32
318  // CHECK: %[[CAST_RESULT:.*]] = "test.cast"(%[[CHANGE_RESULT]]) : (i32) -> i16
319  // CHECK: return %[[CAST_RESULT]]
320  %res = "test.conversion_call_op"(%arg0) { callee=@handle_attr_callee_fn } : (i16) -> (i16)
321  return %res : i16
322}
323
324// Check a function with complex ops is inlined.
325func.func @double_square_complex(%cplx: complex<f32>) -> complex<f32> {
326  %double = complex.add %cplx, %cplx : complex<f32>
327  %square = complex.mul %double, %double : complex<f32>
328  return %square : complex<f32>
329}
330
331// CHECK-LABEL: func @inline_with_complex_ops
332func.func @inline_with_complex_ops() -> complex<f32> {
333  %c1 = arith.constant 1.0 : f32
334  %c2 = arith.constant 2.0 : f32
335  %c = complex.create %c1, %c2 : complex<f32>
336
337  // CHECK: complex.add
338  // CHECK: complex.mul
339  // CHECK-NOT: call
340  %r = call @double_square_complex(%c) : (complex<f32>) -> (complex<f32>)
341  return %r : complex<f32>
342}
343