1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
|
// RUN: mlir-opt %s -linalg-detensorize=aggressive-mode | FileCheck %s -check-prefix=DET-ALL
// RUN: mlir-opt %s -linalg-detensorize | FileCheck %s -check-prefix=DET-CF
#map0 = affine_map<() -> ()>
#map1 = affine_map<(i) -> ()>
#map2 = affine_map<(i) -> (i)>
#attrs = {
indexing_maps = [#map0, #map0, #map0],
iterator_types = []
}
#sum_reduction_attrs = {
indexing_maps = [#map2, #map1],
iterator_types = ["reduction"]
}
#broadcast_attrs = {
indexing_maps = [#map1, #map2],
iterator_types = ["parallel"]
}
func @main(%farg0: tensor<10xi32>, %farg1: tensor<i32>) -> tensor<i32> attributes {} {
br ^bb1(%farg0 : tensor<10xi32>)
^bb1(%0: tensor<10xi32>): // 2 preds: ^bb0, ^bb2
%1 = linalg.init_tensor [] : tensor<i32>
%2 = linalg.generic #sum_reduction_attrs
ins(%0: tensor<10xi32>)
outs(%1: tensor<i32>) {
^bb(%a: i32, %x: i32):
%b = addi %x, %a : i32
linalg.yield %b : i32
} -> tensor<i32>
%3 = linalg.init_tensor [] : tensor<i1>
%4 = linalg.generic #attrs
ins(%2, %farg1 : tensor<i32>, tensor<i32>)
outs(%3 : tensor<i1>) {
^bb0(%arg0: i32, %arg1: i32, %arg2: i1): // no predecessors
%8 = cmpi slt, %arg0, %arg1 : i32
linalg.yield %8 : i1
} -> tensor<i1>
%5 = tensor.extract %4[] : tensor<i1>
cond_br %5, ^bb2(%2 : tensor<i32>), ^bb3(%2 : tensor<i32>)
^bb2(%6: tensor<i32>): // pred: ^bb1
%7 = linalg.init_tensor [10] : tensor<10xi32>
%9 = linalg.generic #broadcast_attrs
ins(%6: tensor<i32>)
outs(%7: tensor<10xi32>) {
^bb(%a: i32, %b: i32) :
linalg.yield %a : i32
} -> tensor<10xi32>
br ^bb1(%9 : tensor<10xi32>)
^bb3(%10: tensor<i32>): // pred: ^bb1
return %10 : tensor<i32>
}
// Test aggresively detensoring all detensorable ops.
//
// DET-ALL-LABEL: func @main
// DET-ALL-SAME: (%{{.*}}: tensor<10xi32>, %{{.*}}: tensor<i32>)
// DET-ALL: br ^[[bb1:.*]](%{{.*}} : tensor<10xi32>)
// DET-ALL: ^[[bb1]](%{{.*}}: tensor<10xi32>)
// DET-ALL: linalg.init_tensor [] : tensor<i32>
// DET-ALL: linalg.generic {{{.*}}} ins(%{{.*}} : tensor<10xi32>) outs(%{{.*}} : tensor<i32>) {
// DET-ALL: ^bb0(%{{.*}}: i32, %{{.*}}: i32): // no predecessors
// DET-ALL: %{{.*}} = addi %{{.*}}, %{{.*}}
// DET-ALL: linalg.yield %{{.*}} : i32
// DET-ALL: } -> tensor<i32>
// DET-ALL: tensor.extract %{{.*}}[] : tensor<i32>
// DET-ALL: tensor.extract %{{.*}}[] : tensor<i32>
// DET-ALL: cmpi slt, %{{.*}}, %{{.*}} : i32
// DET-ALL: tensor.extract %{{.*}}[] : tensor<i32>
// DET-ALL: tensor.extract %{{.*}}[] : tensor<i32>
// DET-ALL: cond_br %{{.*}}, ^[[bb2:.*]](%{{.*}} : i32), ^[[bb3:.*]](%{{.*}} : i32)
// DET-ALL: ^[[bb2]](%{{.*}}: i32)
// DET-ALL: tensor.from_elements %{{.*}} : tensor<1xi32>
// DET-ALL: linalg.tensor_collapse_shape %{{.*}} [] : tensor<1xi32> into tensor<i32>
// DET-ALL: linalg.init_tensor [10] : tensor<10xi32>
// DET-ALL: linalg.generic {{{.*}}} ins(%{{.*}} : tensor<i32>) outs(%{{.*}} : tensor<10xi32>) {
// DET-ALL: ^bb0(%{{.*}}: i32, %{{.*}}: i32):
// DET-ALL: linalg.yield %{{.*}} : i32
// DET-ALL: } -> tensor<10xi32>
// DET-ALL: br ^[[bb1]](%{{.*}} : tensor<10xi32>)
// DET-ALL: ^[[bb3]](%{{.*}}: i32)
// DET-ALL: tensor.from_elements %{{.*}} : tensor<1xi32>
// DET-ALL: linalg.tensor_collapse_shape %{{.*}} [] : tensor<1xi32> into tensor<i32>
// DET-ALL: return %{{.*}} : tensor<i32>
// DET-ALL: }
// Try to detensor pure control-flow. However, that fails since the potential
// detensorable component contains some ops that cannot be detensored.
//
// DET-CF-LABEL: func @main
// DET-CF-SAME: (%{{.*}}: tensor<10xi32>, %{{.*}}: tensor<i32>)
// DET-CF: br ^[[bb1:.*]](%{{.*}} : tensor<10xi32>)
// DET-CF: ^bb1(%{{.*}}: tensor<10xi32>)
// DET-CF: %{{.*}} = linalg.generic {{{.*}}} ins(%{{.*}} : tensor<10xi32>) outs(%{{.*}} : tensor<i32>) {
// DET-CF: %{{.*}} = linalg.generic {{{.*}}} ins(%{{.*}}, %{{.*}} : tensor<i32>, tensor<i32>) outs(%{{.*}} : tensor<i1>) {
// DET-CF: cond_br %{{.*}}, ^bb2(%{{.*}} : tensor<i32>), ^bb3(%{{.*}} : tensor<i32>)
// DET-CF: ^bb2(%{{.*}}: tensor<i32>)
// DET-CF: %{{.*}} = linalg.generic {{{.*}}} ins(%{{.*}} : tensor<i32>) outs(%{{.*}} : tensor<10xi32>) {
// DET-CF: br ^bb1(%{{.*}} : tensor<10xi32>)
// DET-CF: ^bb3(%{{.*}}: tensor<i32>)
// DET-CF: return %{{.*}} : tensor<i32>
// DET-CF: }
|
