canonicalize.mlir
6.36 KB
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
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
// RUN: mlir-opt %s -canonicalize -split-input-file | FileCheck %s
// CHECK-LABEL: func @memref_cast(
func @memref_cast(%a: index, %b: index) -> memref<?x?xf32> {
%c0 = constant 0 : index
%c1 = constant 1 : index
%c8 = constant 8 : index
%c16 = constant 16 : index
%1 = alloc (%b) : memref<?xi8>
%2 = view %1[%c0][] : memref<?xi8> to memref<16x16xf32>
%3 = memref_cast %2 : memref<16x16xf32> to memref<?x?xf32>
%r0 = linalg.range %c0:%c8:%c1 : !linalg.range
// CHECK: linalg.slice {{.*}} : memref<16x16xf32>, !linalg.range, !linalg.range, memref<?x?xf32>
%4 = linalg.slice %3[%r0, %r0] : memref<?x?xf32>, !linalg.range, !linalg.range, memref<?x?xf32>
// CHECK: linalg.matmul{{.*}}: (memref<16x16xf32>, memref<16x16xf32>, memref<16x16xf32>)
linalg.matmul %3, %3, %3 : (memref<?x?xf32>, memref<?x?xf32>, memref<?x?xf32>)
return %4: memref<?x?xf32>
}
// -----
func @collapsing_tensor_reshapes(%arg0 : tensor<?x?x?x?x?xf32>) -> tensor<?x?xf32>
{
%0 = linalg.tensor_reshape %arg0
[affine_map<(d0, d1, d2, d3, d4) -> (d0, d1)>,
affine_map<(d0, d1, d2, d3, d4) -> (d2)>,
affine_map<(d0, d1, d2, d3, d4) -> (d3, d4)>] :
tensor<?x?x?x?x?xf32> into tensor<?x?x?xf32>
%1 = linalg.tensor_reshape %0
[affine_map<(d0, d1, d2) -> (d0, d1)>,
affine_map<(d0, d1, d2) -> (d2)>] :
tensor<?x?x?xf32> into tensor<?x?xf32>
return %1 : tensor<?x?xf32>
}
// CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2)>
// CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1, d2, d3, d4) -> (d3, d4)>
// CHECK-LABEL: collapsing_tensor_reshapes
// CHECK: linalg.tensor_reshape %{{.*}} [#[[$MAP0]], #[[$MAP1]]]
// CHECK-NOT: linalg.tensor_reshape
// -----
func @expanding_tensor_reshapes(%arg0 : tensor<?x?xf32>) -> tensor<?x?x?x?x?xf32>
{
%0 = linalg.tensor_reshape %arg0
[affine_map<(d0, d1, d2) -> (d0, d1)>,
affine_map<(d0, d1, d2) -> (d2)>] :
tensor<?x?xf32> into tensor<?x?x?xf32>
%1 = linalg.tensor_reshape %0
[affine_map<(d0, d1, d2, d3, d4) -> (d0, d1)>,
affine_map<(d0, d1, d2, d3, d4) -> (d2)>,
affine_map<(d0, d1, d2, d3, d4) -> (d3, d4)>] :
tensor<?x?x?xf32> into tensor<?x?x?x?x?xf32>
return %1 : tensor<?x?x?x?x?xf32>
}
// CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2)>
// CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1, d2, d3, d4) -> (d3, d4)>
// CHECK-LABEL: expanding_tensor_reshapes
// CHECK: linalg.tensor_reshape %{{.*}} [#[[$MAP0]], #[[$MAP1]]]
// CHECK-NOT: linalg.tensor_reshape
// -----
func @collapsing_memref_reshapes(%arg0 : memref<?x?x?x?x?xf32>) -> memref<?x?xf32>
{
%0 = linalg.reshape %arg0
[affine_map<(d0, d1, d2, d3, d4) -> (d0, d1)>,
affine_map<(d0, d1, d2, d3, d4) -> (d2)>,
affine_map<(d0, d1, d2, d3, d4) -> (d3, d4)>] :
memref<?x?x?x?x?xf32> into memref<?x?x?xf32>
%1 = linalg.reshape %0
[affine_map<(d0, d1, d2) -> (d0, d1)>,
affine_map<(d0, d1, d2) -> (d2)>] :
memref<?x?x?xf32> into memref<?x?xf32>
return %1 : memref<?x?xf32>
}
// CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2)>
// CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1, d2, d3, d4) -> (d3, d4)>
// CHECK-LABEL: collapsing_memref_reshapes
// CHECK: linalg.reshape %{{.*}} [#[[$MAP0]], #[[$MAP1]]]
// CHECK-NOT: linalg.reshape
// -----
func @expanding_memref_reshapes(%arg0 : memref<?x?xf32>) -> memref<?x?x?x?x?xf32>
{
%0 = linalg.reshape %arg0
[affine_map<(d0, d1, d2) -> (d0, d1)>,
affine_map<(d0, d1, d2) -> (d2)>] :
memref<?x?xf32> into memref<?x?x?xf32>
%1 = linalg.reshape %0
[affine_map<(d0, d1, d2, d3, d4) -> (d0, d1)>,
affine_map<(d0, d1, d2, d3, d4) -> (d2)>,
affine_map<(d0, d1, d2, d3, d4) -> (d3, d4)>] :
memref<?x?x?xf32> into memref<?x?x?x?x?xf32>
return %1 : memref<?x?x?x?x?xf32>
}
// CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2)>
// CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1, d2, d3, d4) -> (d3, d4)>
// CHECK-LABEL: expanding_memref_reshapes
// CHECK: linalg.reshape %{{.*}} [#[[$MAP0]], #[[$MAP1]]]
// CHECK-NOT: linalg.reshape
// -----
func @fold_tensor_reshape(%arg0 : tensor<12x4xf32>) -> tensor<12x4xf32>
{
%0 = linalg.tensor_reshape %arg0
[affine_map<(d0, d1, d2) -> (d0, d1)>,
affine_map<(d0, d1, d2) -> (d2)>] :
tensor<12x4xf32> into tensor<3x4x4xf32>
%1 = linalg.tensor_reshape %0
[affine_map<(d0, d1, d2) -> (d0, d1)>,
affine_map<(d0, d1, d2) -> (d2)>] :
tensor<3x4x4xf32> into tensor<12x4xf32>
return %1 : tensor<12x4xf32>
}
// CHECK-LABEL: @fold_tensor_reshape
// CHECK-NOT: linalg.tensor_reshape
// -----
func @no_fold_tensor_reshape(%arg0 : tensor<?x?xf32>) -> tensor<?x?xf32>
{
%0 = linalg.tensor_reshape %arg0
[affine_map<(d0, d1, d2) -> (d0, d1)>,
affine_map<(d0, d1, d2) -> (d2)>] :
tensor<?x?xf32> into tensor<?x?x?xf32>
%1 = linalg.tensor_reshape %0
[affine_map<(d0, d1, d2) -> (d0, d1)>,
affine_map<(d0, d1, d2) -> (d2)>] :
tensor<?x?x?xf32> into tensor<?x?xf32>
return %1 : tensor<?x?xf32>
}
// CHECK-LABEL: @no_fold_tensor_reshape
// CHECK: linalg.tensor_reshape
// CHECK: linalg.tensor_reshape
// -----
func @fold_memref_reshape(%arg0 : memref<12x4xf32>) -> memref<12x4xf32>
{
%0 = linalg.reshape %arg0
[affine_map<(d0, d1, d2) -> (d0, d1)>,
affine_map<(d0, d1, d2) -> (d2)>] :
memref<12x4xf32> into memref<3x4x4xf32>
%1 = linalg.reshape %0
[affine_map<(d0, d1, d2) -> (d0, d1)>,
affine_map<(d0, d1, d2) -> (d2)>] :
memref<3x4x4xf32> into memref<12x4xf32>
return %1 : memref<12x4xf32>
}
// CHECK-LABEL: @fold_memref_reshape
// CHECK-NOT: linalg.reshape
// -----
func @no_fold_memref_reshape(%arg0 : memref<?x?xf32>) -> memref<?x?xf32>
{
%0 = linalg.reshape %arg0
[affine_map<(d0, d1, d2) -> (d0, d1)>,
affine_map<(d0, d1, d2) -> (d2)>] :
memref<?x?xf32> into memref<?x?x?xf32>
%1 = linalg.reshape %0
[affine_map<(d0, d1, d2) -> (d0, d1)>,
affine_map<(d0, d1, d2) -> (d2)>] :
memref<?x?x?xf32> into memref<?x?xf32>
return %1 : memref<?x?xf32>
}
// CHECK-LABEL: @no_fold_memref_reshape
// CHECK: linalg.reshape
// CHECK: linalg.reshape