mlir/test/Transforms/inlining.mlir

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