1// RUN: mlir-opt %s -split-input-file -allow-unregistered-dialect -pass-pipeline="func.func(linalg-detensorize)" | FileCheck %s
2
3#map0 = affine_map<() -> ()>
4
5#attrs = {
6  indexing_maps = [#map0, #map0, #map0],
7  iterator_types = []
8}
9
10func.func @main() -> (tensor<i32>) attributes {} {
11  %c0 = arith.constant 0 : i32
12  %0 = tensor.from_elements %c0 : tensor<i32>
13  %c10 = arith.constant 10 : i32
14  %1 = tensor.from_elements %c10 : tensor<i32>
15  cf.br ^bb1(%0 : tensor<i32>)
16
17^bb1(%2: tensor<i32>):  // 2 preds: ^bb0, ^bb2
18  %3 = linalg.init_tensor [] : tensor<i1>
19  %4 = linalg.generic #attrs
20    ins(%2, %1 : tensor<i32>, tensor<i32>)
21    outs(%3 : tensor<i1>) {
22    ^bb0(%arg0: i32, %arg1: i32, %arg2: i1):
23      %8 = arith.cmpi slt, %arg0, %arg1 : i32
24      linalg.yield %8 : i1
25  } -> tensor<i1>
26  %5 = tensor.extract %4[] : tensor<i1>
27  cf.cond_br %5, ^bb2(%2 : tensor<i32>), ^bb3(%2 : tensor<i32>)
28
29^bb2(%6: tensor<i32>):  // pred: ^bb1
30  %7 = linalg.init_tensor [] : tensor<i32>
31  %8 = linalg.generic #attrs
32    ins(%6, %6 : tensor<i32>, tensor<i32>)
33    outs(%7 : tensor<i32>) {
34    ^bb0(%arg0: i32, %arg1: i32, %arg2: i32):
35      %9 = arith.addi %arg0, %arg1 : i32
36      linalg.yield %9 : i32
37  } -> tensor<i32>
38  cf.br ^bb3(%8 : tensor<i32>)
39
40^bb3(%10: tensor<i32>):  // pred: ^bb1
41  return %10 : tensor<i32>
42}
43
44// CHECK-LABEL:  func @main()
45// CHECK-DAG:     arith.constant 0
46// CHECK-DAG:     arith.constant 10
47// CHECK:         cf.br ^[[bb1:.*]](%{{.*}}: i32)
48// CHECK-NEXT:   ^[[bb1]](%{{.*}}: i32):
49// CHECK-NEXT:     arith.cmpi slt, %{{.*}}, %{{.*}}
50// CHECK-NEXT:     cf.cond_br %{{.*}}, ^[[bb2:.*]](%{{.*}} : i32), ^bb3(%{{.*}} : i32)
51// CHECK-NEXT:   ^[[bb2]](%{{.*}}: i32)
52// CHECK-NEXT:     arith.addi %{{.*}}, %{{.*}}
53// CHECK-NEXT:     cf.br ^[[bb3:.*]](%{{.*}} : i32)
54// CHECK-NEXT:   ^[[bb3]](%{{.*}}: i32)
55// CHECK-NEXT:     tensor.from_elements %{{.*}} : tensor<i32>
56// CHECK-NEXT:     return %{{.*}}
57// CHECK-NEXT:   }
58
59// -----
60
61// Similar to the above test with one change: one of the block after the
62// if-condition passes/forwards its tensor argument to another block.
63
64#map0 = affine_map<() -> ()>
65
66#attrs = {
67  indexing_maps = [#map0, #map0, #map0],
68  iterator_types = []
69}
70
71func.func @main() -> (tensor<i32>) attributes {} {
72  %c0 = arith.constant 0 : i32
73  %0 = tensor.from_elements %c0 : tensor<i32>
74  %c10 = arith.constant 10 : i32
75  %1 = tensor.from_elements %c10 : tensor<i32>
76  cf.br ^bb1(%0 : tensor<i32>)
77
78^bb1(%2: tensor<i32>):  // 2 preds: ^bb0, ^bb2
79  %3 = linalg.init_tensor [] : tensor<i1>
80  %4 = linalg.generic #attrs
81    ins(%2, %1 : tensor<i32>, tensor<i32>)
82    outs(%3 : tensor<i1>) {
83    ^bb0(%arg0: i32, %arg1: i32, %arg2: i1):
84      %8 = arith.cmpi slt, %arg0, %arg1 : i32
85      linalg.yield %8 : i1
86  } -> tensor<i1>
87  %5 = tensor.extract %4[] : tensor<i1>
88  cf.cond_br %5, ^bb2(%2 : tensor<i32>), ^bb3(%2 : tensor<i32>)
89
90^bb2(%6: tensor<i32>):  // pred: ^bb1
91  %7 = linalg.init_tensor [] : tensor<i32>
92  %8 = linalg.generic #attrs
93    ins(%6, %6 : tensor<i32>, tensor<i32>)
94    outs(%7 : tensor<i32>) {
95    ^bb0(%arg0: i32, %arg1: i32, %arg2: i32):
96      %9 = arith.addi %arg0, %arg1 : i32
97      linalg.yield %9 : i32
98  } -> tensor<i32>
99  cf.br ^bb3(%8 : tensor<i32>)
100
101^bb3(%10: tensor<i32>):  // pred: ^bb1
102  cf.br ^bb4(%10 : tensor<i32>)
103
104^bb4(%11: tensor<i32>):  // pred: ^bb1
105  return %11 : tensor<i32>
106}
107
108// CHECK-LABEL:  func @main()
109// CHECK-DAG:     arith.constant 0
110// CHECK-DAG:     arith.constant 10
111// CHECK:         cf.br ^[[bb1:.*]](%{{.*}}: i32)
112// CHECK-NEXT:   ^[[bb1]](%{{.*}}: i32):
113// CHECK-NEXT:     arith.cmpi slt, %{{.*}}, %{{.*}}
114// CHECK-NEXT:     cf.cond_br %{{.*}}, ^[[bb2:.*]](%{{.*}} : i32), ^bb3(%{{.*}} : i32)
115// CHECK-NEXT:   ^[[bb2]](%{{.*}}: i32)
116// CHECK-NEXT:     arith.addi %{{.*}}, %{{.*}}
117// CHECK-NEXT:     cf.br ^[[bb3:.*]](%{{.*}} : i32)
118// CHECK-NEXT:   ^[[bb3]](%{{.*}}: i32)
119// CHECK-NEXT:     cf.br ^[[bb4:.*]](%{{.*}} : i32)
120// CHECK-NEXT:   ^[[bb4]](%{{.*}}: i32)
121// CHECK-NEXT:     tensor.from_elements %{{.*}} : tensor<i32>
122// CHECK-NEXT:     return %{{.*}}
123// CHECK-NEXT:   }
124
125// -----
126
127#map0 = affine_map<() -> ()>
128
129#attrs = {
130  indexing_maps = [#map0, #map0, #map0],
131  iterator_types = []
132}
133
134func.func @main() -> (tensor<i32>) attributes {} {
135  %c0 = arith.constant 0 : i32
136  %0 = tensor.from_elements %c0 : tensor<i32>
137  %c10 = arith.constant 10 : i32
138  %1 = tensor.from_elements %c10 : tensor<i32>
139  cf.br ^bb1(%0 : tensor<i32>)
140
141^bb1(%2: tensor<i32>):  // 2 preds: ^bb0, ^bb2
142  %3 = linalg.init_tensor [] : tensor<i1>
143  %4 = linalg.generic #attrs
144    ins(%2, %1 : tensor<i32>, tensor<i32>)
145    outs(%3 : tensor<i1>) {
146    ^bb0(%arg0: i32, %arg1: i32, %arg2: i1):
147      %8 = arith.cmpi slt, %arg0, %arg1 : i32
148      linalg.yield %8 : i1
149  } -> tensor<i1>
150  %5 = tensor.extract %4[] : tensor<i1>
151  // This cf.cond_br intentionally has bb2 as it's target for both branches. This
152  // is to make sure that the "forward phase" of the cost-model correctly adds
153  // the users of a block argument (in this case bb2's argument) to the work
154  // list.
155  cf.cond_br %5, ^bb2(%2 : tensor<i32>), ^bb2(%2 : tensor<i32>)
156
157^bb2(%6: tensor<i32>):  // pred: ^bb1
158  %12 = tensor.from_elements %c10 : tensor<i32>
159  %7 = linalg.init_tensor [] : tensor<i32>
160  %8 = linalg.generic #attrs
161    ins(%6, %12 : tensor<i32>, tensor<i32>)
162    outs(%7 : tensor<i32>) {
163    ^bb0(%arg0: i32, %arg1: i32, %arg2: i32):
164      %9 = arith.addi %arg0, %arg1 : i32
165      linalg.yield %9 : i32
166  } -> tensor<i32>
167  cf.br ^bb3(%8 : tensor<i32>)
168
169^bb3(%10: tensor<i32>):  // pred: ^bb1
170  return %10 : tensor<i32>
171}
172
173// CHECK-LABEL:  func @main()
174// CHECK-DAG:     arith.constant 0
175// CHECK-DAG:     arith.constant 10
176// CHECK:         cf.br ^[[bb1:.*]](%{{.*}}: i32)
177// CHECK-NEXT:   ^[[bb1]](%{{.*}}: i32):
178// CHECK-NEXT:     arith.cmpi slt, %{{.*}}, %{{.*}}
179// CHECK-NEXT:     cf.cond_br %{{.*}}, ^[[bb2:.*]](%{{.*}} : i32), ^bb2(%{{.*}} : i32)
180// CHECK-NEXT:   ^[[bb2]](%{{.*}}: i32)
181// CHECK-NEXT:     arith.addi %{{.*}}, %{{.*}}
182// CHECK-NEXT:     cf.br ^[[bb3:.*]](%{{.*}} : i32)
183// CHECK-NEXT:   ^[[bb3]](%{{.*}}: i32)
184// CHECK-NEXT:     tensor.from_elements %{{.*}} : tensor<i32>
185// CHECK-NEXT:     return %{{.*}}
186// CHECK-NEXT:   }
187