xref: /llvm-project/mlir/test/Dialect/MemRef/mem2reg.mlir (revision ead8e9d7953e817c52fdfaf7196dfeb2199dab26)
1// RUN: mlir-opt %s --pass-pipeline='builtin.module(func.func(mem2reg{region-simplify=false}))' --split-input-file | FileCheck %s
2
3// CHECK-LABEL: func.func @basic
4func.func @basic() -> i32 {
5  // CHECK-NOT: = memref.alloca
6  // CHECK: %[[RES:.*]] = arith.constant 5 : i32
7  // CHECK-NOT: = memref.alloca
8  %0 = arith.constant 5 : i32
9  %1 = memref.alloca() : memref<i32>
10  memref.store %0, %1[] : memref<i32>
11  %2 = memref.load %1[] : memref<i32>
12  // CHECK: return %[[RES]] : i32
13  return %2 : i32
14}
15
16// -----
17
18// CHECK-LABEL: func.func @basic_default
19func.func @basic_default() -> i32 {
20  // CHECK-NOT: = memref.alloca
21  // CHECK: %[[RES:.*]] = arith.constant 0 : i32
22  // CHECK-NOT: = memref.alloca
23  %0 = arith.constant 5 : i32
24  %1 = memref.alloca() : memref<i32>
25  %2 = memref.load %1[] : memref<i32>
26  // CHECK-NOT: memref.store
27  memref.store %0, %1[] : memref<i32>
28  // CHECK: return %[[RES]] : i32
29  return %2 : i32
30}
31
32// -----
33
34// CHECK-LABEL: func.func @basic_float
35func.func @basic_float() -> f32 {
36  // CHECK-NOT: = memref.alloca
37  // CHECK: %[[RES:.*]] = arith.constant {{.*}} : f32
38  %0 = arith.constant 5.2 : f32
39  // CHECK-NOT: = memref.alloca
40  %1 = memref.alloca() : memref<f32>
41  memref.store %0, %1[] : memref<f32>
42  %2 = memref.load %1[] : memref<f32>
43  // CHECK: return %[[RES]] : f32
44  return %2 : f32
45}
46
47// -----
48
49// CHECK-LABEL: func.func @basic_ranked
50func.func @basic_ranked() -> i32 {
51  // CHECK-NOT: = memref.alloca
52  // CHECK: %[[RES:.*]] = arith.constant 5 : i32
53  // CHECK-NOT: = memref.alloca
54  %0 = arith.constant 0 : index
55  %1 = arith.constant 5 : i32
56  %2 = memref.alloca() : memref<1x1xi32>
57  memref.store %1, %2[%0, %0] : memref<1x1xi32>
58  %3 = memref.load %2[%0, %0] : memref<1x1xi32>
59  // CHECK: return %[[RES]] : i32
60  return %3 : i32
61}
62
63// -----
64
65// CHECK-LABEL: func.func @reject_multiple_elements
66func.func @reject_multiple_elements() -> i32 {
67  // CHECK: %[[INDEX:.*]] = arith.constant 0 : index
68  %0 = arith.constant 0 : index
69  // CHECK: %[[STORED:.*]] = arith.constant 5 : i32
70  %1 = arith.constant 5 : i32
71  // CHECK: %[[ALLOCA:.*]] = memref.alloca()
72  %2 = memref.alloca() : memref<1x2xi32>
73  // CHECK: memref.store %[[STORED]], %[[ALLOCA]][%[[INDEX]], %[[INDEX]]]
74  memref.store %1, %2[%0, %0] : memref<1x2xi32>
75  // CHECK: %[[RES:.*]] = memref.load %[[ALLOCA]][%[[INDEX]], %[[INDEX]]]
76  %3 = memref.load %2[%0, %0] : memref<1x2xi32>
77  // CHECK: return %[[RES]] : i32
78  return %3 : i32
79}
80
81// -----
82
83// CHECK-LABEL: func.func @cycle
84// CHECK-SAME: (%[[ARG0:.*]]: i64, %[[ARG1:.*]]: i1, %[[ARG2:.*]]: i64)
85func.func @cycle(%arg0: i64, %arg1: i1, %arg2: i64) {
86  // CHECK-NOT: = memref.alloca
87  %alloca = memref.alloca() : memref<i64>
88  memref.store %arg2, %alloca[] : memref<i64>
89  // CHECK: cf.cond_br %[[ARG1:.*]], ^[[BB1:.*]](%[[ARG2]] : i64), ^[[BB2:.*]](%[[ARG2]] : i64)
90  cf.cond_br %arg1, ^bb1, ^bb2
91// CHECK: ^[[BB1]](%[[USE:.*]]: i64):
92^bb1:
93  %use = memref.load %alloca[] : memref<i64>
94  // CHECK: call @use(%[[USE]])
95  func.call @use(%use) : (i64) -> ()
96  memref.store %arg0, %alloca[] : memref<i64>
97  // CHECK: cf.br ^[[BB2]](%[[ARG0]] : i64)
98  cf.br ^bb2
99// CHECK: ^[[BB2]](%[[FWD:.*]]: i64):
100^bb2:
101  // CHECK: cf.br ^[[BB1]](%[[FWD]] : i64)
102  cf.br ^bb1
103}
104
105func.func @use(%arg: i64) { return }
106
107// -----
108
109// CHECK-LABEL: func.func @recursive
110// CHECK-SAME: (%[[ARG:.*]]: i64)
111func.func @recursive(%arg: i64) -> i64 {
112  // CHECK-NOT: = memref.alloca()
113  %alloca0 = memref.alloca() : memref<memref<memref<i64>>>
114  %alloca1 = memref.alloca() : memref<memref<i64>>
115  %alloca2 = memref.alloca() : memref<i64>
116  memref.store %arg, %alloca2[] : memref<i64>
117  memref.store %alloca2, %alloca1[] : memref<memref<i64>>
118  memref.store %alloca1, %alloca0[] : memref<memref<memref<i64>>>
119  %load0 = memref.load %alloca0[] : memref<memref<memref<i64>>>
120  %load1 = memref.load %load0[] : memref<memref<i64>>
121  %load2 = memref.load %load1[] : memref<i64>
122  // CHECK: return %[[ARG]] : i64
123  return %load2 : i64
124}
125
126// -----
127
128// CHECK-LABEL: func.func @deny_store_of_alloca
129// CHECK-SAME: (%[[ARG:.*]]: memref<memref<i32>>)
130func.func @deny_store_of_alloca(%arg: memref<memref<i32>>) -> i32 {
131  // CHECK: %[[VALUE:.*]] = arith.constant 5 : i32
132  %0 = arith.constant 5 : i32
133  // CHECK: %[[ALLOCA:.*]] = memref.alloca
134  %1 = memref.alloca() : memref<i32>
135  // Storing into the memref is allowed.
136  // CHECK: memref.store %[[VALUE]], %[[ALLOCA]][]
137  memref.store %0, %1[] : memref<i32>
138  // Storing the memref itself is NOT allowed.
139  // CHECK: memref.store %[[ALLOCA]], %[[ARG]][]
140  memref.store %1, %arg[] : memref<memref<i32>>
141  // CHECK: %[[RES:.*]] = memref.load %[[ALLOCA]][]
142  %2 = memref.load %1[] : memref<i32>
143  // CHECK: return %[[RES]] : i32
144  return %2 : i32
145}
146
147// -----
148
149// CHECK-LABEL: func.func @promotable_nonpromotable_intertwined
150func.func @promotable_nonpromotable_intertwined() -> i32 {
151  // CHECK: %[[NON_PROMOTED:.*]] = memref.alloca() : memref<i32>
152  %0 = memref.alloca() : memref<i32>
153  // CHECK-NOT: = memref.alloca() : memref<memref<i32>>
154  %1 = memref.alloca() : memref<memref<i32>>
155  memref.store %0, %1[] : memref<memref<i32>>
156  %2 = memref.load %1[] : memref<memref<i32>>
157  // CHECK: call @use(%[[NON_PROMOTED]])
158  call @use(%0) : (memref<i32>) -> ()
159  // CHECK: %[[RES:.*]] = memref.load %[[NON_PROMOTED]][]
160  %3 = memref.load %0[] : memref<i32>
161  // CHECK: return %[[RES]] : i32
162  return %3 : i32
163}
164
165func.func @use(%arg: memref<i32>) { return }
166