1// RUN: mlir-opt %s -linalg-detensorize=aggressive-mode | FileCheck %s -check-prefix=DET-ALL
2// RUN: mlir-opt %s -linalg-detensorize | FileCheck %s -check-prefix=DET-CF
3
4#map0 = affine_map<() -> ()>
5
6#attrs = {
7  indexing_maps = [#map0, #map0, #map0],
8  iterator_types = []
9}
10
11func @main(%farg0: tensor<i32>, %farg1: tensor<i32>) -> tensor<i32> attributes {} {
12  br ^bb1(%farg0 : tensor<i32>)
13
14^bb1(%0: tensor<i32>):  // 2 preds: ^bb0, ^bb2
15  %1 = linalg.init_tensor [] : tensor<i1>
16  %2 = linalg.generic #attrs
17    ins(%0, %farg1 : tensor<i32>, tensor<i32>)
18    outs(%1 : tensor<i1>) {
19    ^bb0(%arg0: i32, %arg1: i32, %arg2: i1):  // no predecessors
20      %8 = cmpi slt, %arg0, %arg1 : i32
21      linalg.yield %8 : i1
22  } -> tensor<i1>
23  %3 = tensor.extract %2[] : tensor<i1>
24  cond_br %3, ^bb2(%0 : tensor<i32>), ^bb3(%0 : tensor<i32>)
25
26^bb2(%4: tensor<i32>):  // pred: ^bb1
27  %5 = linalg.init_tensor [] : tensor<i32>
28  %6 = linalg.generic #attrs
29    ins(%4, %4 : tensor<i32>, tensor<i32>)
30    outs(%5 : tensor<i32>) {
31    ^bb0(%arg0: i32, %arg1: i32, %arg2: i32):  // no predecessors
32      %8 = addi %arg0, %arg1 : i32
33      linalg.yield %8 : i32
34  } -> tensor<i32>
35  br ^bb1(%6 : tensor<i32>)
36
37^bb3(%7: tensor<i32>):  // pred: ^bb1
38  return %7 : tensor<i32>
39}
40
41// Test aggresively detensoring all detensorable ops.
42//
43// DET-ALL-LABEL: func @main
44// DET-ALL-SAME:    (%{{.*}}: tensor<i32>, %{{.*}}: tensor<i32>)
45// DET-ALL:         tensor.extract {{.*}}
46// DET-ALL:         br ^[[bb1:.*]](%{{.*}} : i32)
47// DET-ALL:       ^[[bb1]](%{{.*}}: i32)
48// DET-ALL:         cmpi slt, {{.*}}
49// DET-ALL:         cond_br {{.*}}, ^[[bb2:.*]](%{{.*}} : i32), ^[[bb3:.*]](%{{.*}} : i32)
50// DET-ALL:       ^[[bb2]](%{{.*}}: i32)
51// DET-ALL:         addi {{.*}}
52// DET-ALL:         br ^[[bb1]](%{{.*}} : i32)
53// DET-ALL:       ^[[bb3]](%{{.*}}: i32)
54// DET-ALL:         tensor.from_elements {{.*}}
55// DET-ALL:         linalg.tensor_collapse_shape {{.*}}
56// DET-ALL:         return %{{.*}} : tensor<i32>
57
58// Test detensoring only ops involed in control-flow.
59//
60// DET-CF-LABEL: func @main
61// DET-CF-SAME:    (%{{.*}}: tensor<i32>, %{{.*}}: tensor<i32>)
62// DET-CF:         tensor.extract {{.*}}
63// DET-CF:         br ^[[bb1:.*]](%{{.*}} : i32)
64// DET-CF:       ^[[bb1]](%{{.*}}: i32)
65// DET-CF:         cmpi slt, {{.*}}
66// DET-CF:         cond_br {{.*}}, ^[[bb2:.*]](%{{.*}} : i32), ^[[bb3:.*]](%{{.*}} : i32)
67// DET-CF:       ^[[bb2]](%{{.*}}: i32)
68// DET-CF:         addi {{.*}}
69// DET-CF:         br ^[[bb1]](%{{.*}} : i32)
70// DET-CF:       ^[[bb3]](%{{.*}}: i32)
71// DET-CF:         tensor.from_elements %{{.*}} : tensor<1xi32>
72// DET-CF:         linalg.tensor_collapse_shape %{{.*}} [] : tensor<1xi32> into tensor<i32>
73// DET-CF:         return %{{.*}} : tensor<i32>
74