fold-affine-min-scf.mlir
5.29 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
// RUN: mlir-opt %s -test-linalg-transform-patterns=test-affine-min-scf-canonicalization-patterns | FileCheck %s
// CHECK-LABEL: scf_for
func @scf_for(%A : memref<i64>, %step : index) {
%c0 = constant 0 : index
%c1 = constant 1 : index
%c2 = constant 2 : index
%c7 = constant 7 : index
%c4 = constant 4 : index
%c16 = constant 16 : index
%c1024 = constant 1024 : index
// CHECK: scf.for
// CHECK-NEXT: %[[C2:.*]] = constant 2 : index
// CHECK-NEXT: %[[C2I64:.*]] = index_cast %[[C2:.*]]
// CHECK-NEXT: store %[[C2I64]], %{{.*}}[] : memref<i64>
scf.for %i = %c0 to %c4 step %c2 {
%1 = affine.min affine_map<(d0, d1)[] -> (2, d1 - d0)> (%i, %c4)
%2 = index_cast %1: index to i64
store %2, %A[]: memref<i64>
}
// CHECK: scf.for
// CHECK-NEXT: %[[C2:.*]] = constant 2 : index
// CHECK-NEXT: %[[C2I64:.*]] = index_cast %[[C2:.*]]
// CHECK-NEXT: store %[[C2I64]], %{{.*}}[] : memref<i64>
scf.for %i = %c1 to %c7 step %c2 {
%1 = affine.min affine_map<(d0)[s0] -> (s0 - d0, 2)> (%i)[%c7]
%2 = index_cast %1: index to i64
store %2, %A[]: memref<i64>
}
// This should not canonicalize because: 4 - %i may take the value 1 < 2.
// CHECK: scf.for
// CHECK: affine.min
// CHECK: index_cast
scf.for %i = %c1 to %c4 step %c2 {
%1 = affine.min affine_map<(d0)[s0] -> (2, s0 - d0)> (%i)[%c4]
%2 = index_cast %1: index to i64
store %2, %A[]: memref<i64>
}
// This should not canonicalize because: 16 - %i may take the value 15 < 1024.
// CHECK: scf.for
// CHECK: affine.min
// CHECK: index_cast
scf.for %i = %c1 to %c16 step %c1024 {
%1 = affine.min affine_map<(d0) -> (1024, 16 - d0)> (%i)
%2 = index_cast %1: index to i64
store %2, %A[]: memref<i64>
}
// This example should simplify but affine_map is currently missing
// semi-affine canonicalizations: `((s0 * 42 - 1) floordiv s0) * s0`
// should evaluate to 41 * s0.
// Note that this may require positivity assumptions on `s0`.
// Revisit when support is added.
// CHECK: scf.for
// CHECK: affine.min
// CHECK: index_cast
%ub = affine.apply affine_map<(d0) -> (42 * d0)> (%step)
scf.for %i = %c0 to %ub step %step {
%1 = affine.min affine_map<(d0, d1, d2) -> (d0, d1 - d2)> (%step, %ub, %i)
%2 = index_cast %1: index to i64
store %2, %A[]: memref<i64>
}
// This example should simplify but affine_map is currently missing
// semi-affine canonicalizations.
// This example should simplify but affine_map is currently missing
// semi-affine canonicalizations: ` -(((s0 * s0 - 1) floordiv s0) * s0)`
// should evaluate to (s0 - 1) * s0.
// Note that this may require positivity assumptions on `s0`.
// Revisit when support is added.
// CHECK: scf.for
// CHECK: affine.min
// CHECK: index_cast
%ub2 = affine.apply affine_map<(d0)[s0] -> (s0 * d0)> (%step)[%step]
scf.for %i = %c0 to %ub2 step %step {
%1 = affine.min affine_map<(d0, d1, d2) -> (d0, d2 - d1)> (%step, %i, %ub2)
%2 = index_cast %1: index to i64
store %2, %A[]: memref<i64>
}
return
}
// CHECK-LABEL: scf_parallel
func @scf_parallel(%A : memref<i64>, %step : index) {
%c0 = constant 0 : index
%c1 = constant 1 : index
%c2 = constant 2 : index
%c7 = constant 7 : index
%c4 = constant 4 : index
// CHECK: scf.parallel
// CHECK-NEXT: %[[C2:.*]] = constant 2 : index
// CHECK-NEXT: %[[C2I64:.*]] = index_cast %[[C2:.*]]
// CHECK-NEXT: store %[[C2I64]], %{{.*}}[] : memref<i64>
scf.parallel (%i) = (%c0) to (%c4) step (%c2) {
%1 = affine.min affine_map<(d0, d1)[] -> (2, d1 - d0)> (%i, %c4)
%2 = index_cast %1: index to i64
store %2, %A[]: memref<i64>
}
// CHECK: scf.parallel
// CHECK-NEXT: %[[C2:.*]] = constant 2 : index
// CHECK-NEXT: %[[C2I64:.*]] = index_cast %[[C2:.*]]
// CHECK-NEXT: store %[[C2I64]], %{{.*}}[] : memref<i64>
scf.parallel (%i) = (%c1) to (%c7) step (%c2) {
%1 = affine.min affine_map<(d0)[s0] -> (2, s0 - d0)> (%i)[%c7]
%2 = index_cast %1: index to i64
store %2, %A[]: memref<i64>
}
// This example should simplify but affine_map is currently missing
// semi-affine canonicalizations.
// This affine map does not currently evaluate to (0, 0):
// (d0)[s0] -> (s0 mod s0, (-((d0 floordiv s0) * s0) + s0 * 42) mod s0)
// TODO: Revisit when support is added.
// CHECK: scf.parallel
// CHECK: affine.min
// CHECK: index_cast
%ub = affine.apply affine_map<(d0) -> (42 * d0)> (%step)
scf.parallel (%i) = (%c0) to (%ub) step (%step) {
%1 = affine.min affine_map<(d0, d1, d2) -> (d0, d2 - d1)> (%step, %i, %ub)
%2 = index_cast %1: index to i64
store %2, %A[]: memref<i64>
}
// This example should simplify but affine_map is currently missing
// semi-affine canonicalizations.
// This affine map does not currently evaluate to (0, 0):
// (d0)[s0] -> (s0 mod s0, (-((d0 floordiv s0) * s0) + s0 * s0) mod s0)
// TODO: Revisit when support is added.
// CHECK: scf.parallel
// CHECK: affine.min
// CHECK: index_cast
%ub2 = affine.apply affine_map<(d0)[s0] -> (s0 * d0)> (%step)[%step]
scf.parallel (%i) = (%c0) to (%ub2) step (%step) {
%1 = affine.min affine_map<(d0, d1, d2) -> (d0, d2 - d1)> (%step, %i, %ub2)
%2 = index_cast %1: index to i64
store %2, %A[]: memref<i64>
}
return
}