1// RUN: mlir-opt -allow-unregistered-dialect %s -inline="default-pipeline=''" | FileCheck %s
2
3// Basic test that functions within affine operations are inlined.
4func @func_with_affine_ops(%N: index) {
5  %c = constant 200 : index
6  affine.for %i = 1 to 10 {
7    affine.if affine_set<(i)[N] : (i - 2 >= 0, 4 - i >= 0)>(%i)[%c]  {
8      %w = affine.apply affine_map<(d0,d1)[s0] -> (d0+d1+s0)> (%i, %i) [%N]
9    }
10  }
11  return
12}
13
14// CHECK-LABEL: func @inline_with_affine_ops
15func @inline_with_affine_ops() {
16  %c = constant 1 : index
17
18  // CHECK: affine.for
19  // CHECK-NEXT: affine.if
20  // CHECK-NEXT: affine.apply
21  // CHECK-NOT: call
22  call @func_with_affine_ops(%c) : (index) -> ()
23  return
24}
25
26// CHECK-LABEL: func @not_inline_in_affine_op
27func @not_inline_in_affine_op() {
28  %c = constant 1 : index
29
30  // CHECK-NOT: affine.if
31  // CHECK: call
32  affine.for %i = 1 to 10 {
33    call @func_with_affine_ops(%c) : (index) -> ()
34  }
35  return
36}
37
38// -----
39
40// Test when an invalid operation is nested in an affine op.
41func @func_with_invalid_nested_op() {
42  affine.for %i = 1 to 10 {
43    "foo.opaque"() : () -> ()
44  }
45  return
46}
47
48// CHECK-LABEL: func @not_inline_invalid_nest_op
49func @not_inline_invalid_nest_op() {
50  // CHECK: call @func_with_invalid_nested_op
51  call @func_with_invalid_nested_op() : () -> ()
52  return
53}
54
55// -----
56
57// Test that calls are inlined into affine structures.
58func @func_noop() {
59  return
60}
61
62// CHECK-LABEL: func @inline_into_affine_ops
63func @inline_into_affine_ops() {
64  // CHECK-NOT: call @func_noop
65  affine.for %i = 1 to 10 {
66    call @func_noop() : () -> ()
67  }
68  return
69}
70
71// -----
72
73// Test that calls with dimension arguments are properly inlined.
74func @func_dim(%arg0: index, %arg1: memref<?xf32>) {
75  affine.load %arg1[%arg0] : memref<?xf32>
76  return
77}
78
79// CHECK-LABEL: @inline_dimension
80// CHECK: (%[[ARG0:.*]]: memref<?xf32>)
81func @inline_dimension(%arg0: memref<?xf32>) {
82  // CHECK: affine.for %[[IV:.*]] =
83  affine.for %i = 1 to 42 {
84    // CHECK-NOT: call @func_dim
85    // CHECK: affine.load %[[ARG0]][%[[IV]]]
86    call @func_dim(%i, %arg0) : (index, memref<?xf32>) -> ()
87  }
88  return
89}
90
91// -----
92
93// Test that calls with vector operations are also inlined.
94func @func_vector_dim(%arg0: index, %arg1: memref<32xf32>) {
95  affine.vector_load %arg1[%arg0] : memref<32xf32>, vector<4xf32>
96  return
97}
98
99// CHECK-LABEL: @inline_dimension_vector
100// CHECK: (%[[ARG0:.*]]: memref<32xf32>)
101func @inline_dimension_vector(%arg0: memref<32xf32>) {
102  // CHECK: affine.for %[[IV:.*]] =
103  affine.for %i = 1 to 42 {
104    // CHECK-NOT: call @func_dim
105    // CHECK: affine.vector_load %[[ARG0]][%[[IV]]]
106    call @func_vector_dim(%i, %arg0) : (index, memref<32xf32>) -> ()
107  }
108  return
109}
110
111// -----
112
113// Test that calls that would result in violation of affine value
114// categorization (top-level value stop being top-level) are not inlined.
115func private @get_index() -> index
116
117func @func_top_level(%arg0: memref<?xf32>) {
118  %0 = call @get_index() : () -> index
119  affine.load %arg0[%0] : memref<?xf32>
120  return
121}
122
123// CHECK-LABEL: @no_inline_not_top_level
124func @no_inline_not_top_level(%arg0: memref<?xf32>) {
125  affine.for %i = 1 to 42 {
126    // CHECK: call @func_top_level
127    call @func_top_level(%arg0) : (memref<?xf32>) -> ()
128  }
129  return
130}
131