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
|
// RUN: mlir-opt %s -split-input-file -pass-pipeline="builtin.module(func.func(linalg-fold-unit-extent-dims{fold-one-trip-loops-only}))" | FileCheck %s
#accesses = [
affine_map<(i, j, k, l, m) -> (i, k, m)>,
affine_map<(i, j, k, l, m) -> (i, k, j, l, m)>
]
#trait = {
iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"],
indexing_maps = #accesses,
library_call = "some_external_func"
}
func.func @drop_one_trip_loops(%arg0 : tensor<?x1x?xf32>, %shape: tensor<?x1x?x1x?xf32>) -> tensor<?x1x?x1x?xf32>
{
%0 = linalg.generic #trait
ins(%arg0 : tensor<?x1x?xf32>)
outs(%shape : tensor<?x1x?x1x?xf32>) {
^bb0(%arg1 : f32, %arg2 : f32) :
linalg.yield %arg1 : f32
} -> tensor<?x1x?x1x?xf32>
return %0 : tensor<?x1x?x1x?xf32>
}
// CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1, d2) -> (d0, 0, d2)>
// CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1, d2) -> (d0, 0, d1, 0, d2)>
// CHECK-LABEL: func @drop_one_trip_loops
// CHECK: linalg.generic
// CHECK-SAME: indexing_maps = [#[[$MAP0]], #[[$MAP1]]]
// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel"]
// -----
#map0 = affine_map<(i, j) -> (i, j)>
#access = [#map0, #map0]
#trait = {
iterator_types = ["parallel", "parallel"],
indexing_maps = #access,
library_call = "some_external_func"
}
func.func @drop_all_loops(%arg0 : tensor<1x1xf32>) -> tensor<1x1xf32>
{
%0 = linalg.generic #trait
ins(%arg0 : tensor<1x1xf32>)
outs(%arg0 : tensor<1x1xf32>) {
^bb0(%arg1: f32, %arg2: f32) :
linalg.yield %arg1 : f32
} -> tensor<1x1xf32>
return %0 : tensor<1x1xf32>
}
// CHECK-DAG: #[[$MAP0:.*]] = affine_map<() -> (0, 0)>
// CHECK-LABEL: func @drop_all_loops
// CHECK: linalg.generic
// CHECK-SAME: indexing_maps = [#[[$MAP0]], #[[$MAP0]]]
// CHECK-SAME: iterator_types = []
// -----
#map0 = affine_map<(i, j) -> (i, j)>
#access = [#map0, #map0]
#trait = {
iterator_types = ["parallel", "parallel"],
indexing_maps = #access,
library_call = "some_external_func"
}
func.func @drop_all_loops(%arg0 : memref<1x1xf32>, %arg1 : memref<1x1xf32>)
{
linalg.generic #trait
ins(%arg0 : memref<1x1xf32>)
outs(%arg1 : memref<1x1xf32>) {
^bb0(%arg2: f32, %arg3 : f32) :
linalg.yield %arg2 : f32
}
return
}
// CHECK-DAG: #[[$MAP0:.*]] = affine_map<() -> (0, 0)>
// CHECK-LABEL: func @drop_all_loops
// CHECK: linalg.generic
// CHECK-SAME: indexing_maps = [#[[$MAP0]], #[[$MAP0]]]
// CHECK-SAME: iterator_types = []
// -----
#accesses = [
affine_map<(d0, d1) -> (d0, d1)>,
affine_map<(d0, d1) -> (d1)>
]
#trait = {
indexing_maps = #accesses,
iterator_types = ["parallel", "parallel"],
library_call = "some_external_fn"
}
func.func @leading_dim_1_canonicalization(%arg0: tensor<1x5xf32>, %shape: tensor<5xf32>) -> tensor<5xf32> {
%0 = linalg.generic #trait
ins(%arg0 : tensor<1x5xf32>)
outs(%shape : tensor<5xf32>) {
^bb0(%arg2: f32, %arg3: f32):
linalg.yield %arg2 : f32
} -> tensor<5xf32>
return %0 : tensor<5xf32>
}
// CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0) -> (0, d0)>
// CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0) -> (d0)>
// CHECK-LABEL: func @leading_dim_1_canonicalization
// CHECK: linalg.generic
// CHECK-SAME: indexing_maps = [#[[$MAP0]], #[[$MAP1]]]
// CHECK-SAME: iterator_types = ["parallel"]
|
