xref: /llvm-project/mlir/test/Dialect/MemRef/sroa.mlir (revision b142501e9267e308b6b6e5a051f327aa1e0ee3b7)
1// RUN: mlir-opt %s --pass-pipeline="builtin.module(func.func(sroa))" --split-input-file | FileCheck %s
2
3// CHECK-LABEL: func.func @basic
4// CHECK-SAME: (%[[ARG0:.*]]: i32, %[[ARG1:.*]]: i32)
5func.func @basic(%arg0: i32, %arg1: i32) -> i32 {
6  %c0 = arith.constant 0 : index
7  %c1 = arith.constant 1 : index
8  // CHECK-NOT: = memref.alloca()
9  // CHECK-COUNT-2: = memref.alloca() : memref<i32>
10  // CHECK-NOT: = memref.alloca()
11  %alloca = memref.alloca() : memref<2xi32>
12  // CHECK: memref.store %[[ARG0]], %[[ALLOCA0:.*]][]
13  memref.store %arg0, %alloca[%c0] : memref<2xi32>
14  // CHECK: memref.store %[[ARG1]], %[[ALLOCA1:.*]][]
15  memref.store %arg1, %alloca[%c1] : memref<2xi32>
16  // CHECK: %[[RES:.*]] = memref.load %[[ALLOCA0]][]
17  %res = memref.load %alloca[%c0] : memref<2xi32>
18  // CHECK: return %[[RES]] : i32
19  return %res : i32
20}
21
22// -----
23
24// CHECK-LABEL: func.func @basic_high_dimensions
25// CHECK-SAME: (%[[ARG0:.*]]: i32, %[[ARG1:.*]]: i32, %[[ARG2:.*]]: i32)
26func.func @basic_high_dimensions(%arg0: i32, %arg1: i32, %arg2: i32) -> i32 {
27  %c0 = arith.constant 0 : index
28  %c1 = arith.constant 1 : index
29  // CHECK-NOT: = memref.alloca()
30  // CHECK-COUNT-3: = memref.alloca() : memref<i32>
31  // CHECK-NOT: = memref.alloca()
32  %alloca = memref.alloca() : memref<2x2xi32>
33  // CHECK: memref.store %[[ARG0]], %[[ALLOCA0:.*]][]
34  memref.store %arg0, %alloca[%c0, %c0] : memref<2x2xi32>
35  // CHECK: memref.store %[[ARG1]], %[[ALLOCA1:.*]][]
36  memref.store %arg1, %alloca[%c0, %c1] : memref<2x2xi32>
37  // CHECK: memref.store %[[ARG2]], %[[ALLOCA2:.*]][]
38  memref.store %arg2, %alloca[%c1, %c0] : memref<2x2xi32>
39  // CHECK: %[[RES:.*]] = memref.load %[[ALLOCA1]][]
40  %res = memref.load %alloca[%c0, %c1] : memref<2x2xi32>
41  // CHECK: return %[[RES]] : i32
42  return %res : i32
43}
44
45// -----
46
47// CHECK-LABEL: func.func @resolve_alias
48// CHECK-SAME: (%[[ARG0:.*]]: i32, %[[ARG1:.*]]: i32)
49func.func @resolve_alias(%arg0: i32, %arg1: i32) -> i32 {
50  %c0 = arith.constant 0 : index
51  // CHECK-NOT: = memref.alloca()
52  // CHECK: %[[ALLOCA:.*]] = memref.alloca() : memref<i32>
53  // CHECK-NOT: = memref.alloca()
54  %alloca = memref.alloca() : memref<2xi32>
55  // CHECK: memref.store %[[ARG0]], %[[ALLOCA]][]
56  memref.store %arg0, %alloca[%c0] : memref<2xi32>
57  // CHECK: memref.store %[[ARG1]], %[[ALLOCA]][]
58  memref.store %arg1, %alloca[%c0] : memref<2xi32>
59  // CHECK: %[[RES:.*]] = memref.load %[[ALLOCA]][]
60  %res = memref.load %alloca[%c0] : memref<2xi32>
61  // CHECK: return %[[RES]] : i32
62  return %res : i32
63}
64
65// -----
66
67// CHECK-LABEL: func.func @no_direct_use
68// CHECK-SAME: (%[[ARG0:.*]]: i32, %[[ARG1:.*]]: i32)
69func.func @no_direct_use(%arg0: i32, %arg1: i32) -> i32 {
70  // CHECK: %[[C0:.*]] = arith.constant 0 : index
71  %c0 = arith.constant 0 : index
72  // CHECK: %[[C1:.*]] = arith.constant 1 : index
73  %c1 = arith.constant 1 : index
74  // CHECK-NOT: = memref.alloca()
75  // CHECK: %[[ALLOCA:.*]] = memref.alloca() : memref<2xi32>
76  // CHECK-NOT: = memref.alloca()
77  %alloca = memref.alloca() : memref<2xi32>
78  // CHECK: memref.store %[[ARG0]], %[[ALLOCA]][%[[C0]]]
79  memref.store %arg0, %alloca[%c0] : memref<2xi32>
80  // CHECK: memref.store %[[ARG1]], %[[ALLOCA]][%[[C1]]]
81  memref.store %arg1, %alloca[%c1] : memref<2xi32>
82  // CHECK: %[[RES:.*]] = memref.load %[[ALLOCA]][%[[C0]]]
83  %res = memref.load %alloca[%c0] : memref<2xi32>
84  call @use(%alloca) : (memref<2xi32>) -> ()
85  // CHECK: return %[[RES]] : i32
86  return %res : i32
87}
88
89func.func @use(%foo: memref<2xi32>) { return }
90
91// -----
92
93// CHECK-LABEL: func.func @no_dynamic_indexing
94// CHECK-SAME: (%[[ARG0:.*]]: i32, %[[ARG1:.*]]: i32, %[[INDEX:.*]]: index)
95func.func @no_dynamic_indexing(%arg0: i32, %arg1: i32, %index: index) -> i32 {
96  // CHECK: %[[C0:.*]] = arith.constant 0 : index
97  %c0 = arith.constant 0 : index
98  // CHECK-NOT: = memref.alloca()
99  // CHECK: %[[ALLOCA:.*]] = memref.alloca() : memref<2xi32>
100  // CHECK-NOT: = memref.alloca()
101  %alloca = memref.alloca() : memref<2xi32>
102  // CHECK: memref.store %[[ARG0]], %[[ALLOCA]][%[[C0]]]
103  memref.store %arg0, %alloca[%c0] : memref<2xi32>
104  // CHECK: memref.store %[[ARG1]], %[[ALLOCA]][%[[INDEX]]]
105  memref.store %arg1, %alloca[%index] : memref<2xi32>
106  // CHECK: %[[RES:.*]] = memref.load %[[ALLOCA]][%[[C0]]]
107  %res = memref.load %alloca[%c0] : memref<2xi32>
108  // CHECK: return %[[RES]] : i32
109  return %res : i32
110}
111
112// -----
113
114// CHECK-LABEL: func.func @no_dynamic_shape
115// CHECK-SAME: (%[[ARG0:.*]]: i32, %[[ARG1:.*]]: i32)
116func.func @no_dynamic_shape(%arg0: i32, %arg1: i32) -> i32 {
117  // CHECK: %[[C0:.*]] = arith.constant 0 : index
118  %c0 = arith.constant 0 : index
119  // CHECK: %[[C1:.*]] = arith.constant 1 : index
120  %c1 = arith.constant 1 : index
121  // CHECK-NOT: = memref.alloca()
122  // CHECK: %[[ALLOCA:.*]] = memref.alloca(%[[C1]]) : memref<?x2xi32>
123  // CHECK-NOT: = memref.alloca()
124  %alloca = memref.alloca(%c1) : memref<?x2xi32>
125  // CHECK: memref.store %[[ARG0]], %[[ALLOCA]][%[[C0]], %[[C0]]]
126  memref.store %arg0, %alloca[%c0, %c0] : memref<?x2xi32>
127  // CHECK: %[[RES:.*]] = memref.load %[[ALLOCA]][%[[C0]], %[[C0]]]
128  %res = memref.load %alloca[%c0, %c0] : memref<?x2xi32>
129  // CHECK: return %[[RES]] : i32
130  return %res : i32
131}
132
133// -----
134
135// CHECK-LABEL: func.func @no_out_of_bound_write
136// CHECK-SAME: (%[[ARG0:.*]]: i32, %[[ARG1:.*]]: i32)
137func.func @no_out_of_bound_write(%arg0: i32, %arg1: i32) -> i32 {
138  // CHECK: %[[C0:.*]] = arith.constant 0 : index
139  %c0 = arith.constant 0 : index
140  // CHECK: %[[C100:.*]] = arith.constant 100 : index
141  %c100 = arith.constant 100 : index
142  // CHECK-NOT: = memref.alloca()
143  // CHECK: %[[ALLOCA:.*]] = memref.alloca() : memref<2xi32>
144  // CHECK-NOT: = memref.alloca()
145  %alloca = memref.alloca() : memref<2xi32>
146  // CHECK: memref.store %[[ARG0]], %[[ALLOCA]][%[[C0]]]
147  memref.store %arg0, %alloca[%c0] : memref<2xi32>
148  // CHECK: memref.store %[[ARG1]], %[[ALLOCA]][%[[C100]]]
149  memref.store %arg1, %alloca[%c100] : memref<2xi32>
150  // CHECK: %[[RES:.*]] = memref.load %[[ALLOCA]][%[[C0]]]
151  %res = memref.load %alloca[%c0] : memref<2xi32>
152  // CHECK: return %[[RES]] : i32
153  return %res : i32
154}
155
156// -----
157
158// CHECK-LABEL: func.func @no_out_of_bound_load
159// CHECK-SAME: (%[[ARG0:.*]]: i32, %[[ARG1:.*]]: i32)
160func.func @no_out_of_bound_load(%arg0: i32, %arg1: i32) -> i32 {
161  // CHECK: %[[C0:.*]] = arith.constant 0 : index
162  %c0 = arith.constant 0 : index
163  // CHECK: %[[C100:.*]] = arith.constant 100 : index
164  %c100 = arith.constant 100 : index
165  // CHECK-NOT: = memref.alloca()
166  // CHECK: %[[ALLOCA:.*]] = memref.alloca() : memref<2xi32>
167  // CHECK-NOT: = memref.alloca()
168  %alloca = memref.alloca() : memref<2xi32>
169  // CHECK: memref.store %[[ARG0]], %[[ALLOCA]][%[[C0]]]
170  memref.store %arg0, %alloca[%c0] : memref<2xi32>
171  // CHECK: %[[RES:.*]] = memref.load %[[ALLOCA]][%[[C100]]]
172  %res = memref.load %alloca[%c100] : memref<2xi32>
173  // CHECK: return %[[RES]] : i32
174  return %res : i32
175}
176