1// RUN: mlir-opt %s -test-linalg-transform-patterns=test-affine-min-scf-canonicalization-patterns | FileCheck %s
2
3// CHECK-LABEL: scf_for
4func @scf_for(%A : memref<i64>, %step : index) {
5  %c0 = constant 0 : index
6  %c1 = constant 1 : index
7  %c2 = constant 2 : index
8  %c7 = constant 7 : index
9  %c4 = constant 4 : index
10  %c16 = constant 16 : index
11  %c1024 = constant 1024 : index
12
13  //      CHECK: scf.for
14  // CHECK-NEXT:   %[[C2:.*]] = constant 2 : index
15  // CHECK-NEXT:   %[[C2I64:.*]] = index_cast %[[C2:.*]]
16  // CHECK-NEXT:   store %[[C2I64]], %{{.*}}[] : memref<i64>
17  scf.for %i = %c0 to %c4 step %c2 {
18    %1 = affine.min affine_map<(d0, d1)[] -> (2, d1 - d0)> (%i, %c4)
19    %2 = index_cast %1: index to i64
20    store %2, %A[]: memref<i64>
21  }
22
23  //      CHECK: scf.for
24  // CHECK-NEXT:   %[[C2:.*]] = constant 2 : index
25  // CHECK-NEXT:   %[[C2I64:.*]] = index_cast %[[C2:.*]]
26  // CHECK-NEXT:   store %[[C2I64]], %{{.*}}[] : memref<i64>
27  scf.for %i = %c1 to %c7 step %c2 {
28    %1 = affine.min affine_map<(d0)[s0] -> (s0 - d0, 2)> (%i)[%c7]
29    %2 = index_cast %1: index to i64
30    store %2, %A[]: memref<i64>
31  }
32
33  // This should not canonicalize because: 4 - %i may take the value 1 < 2.
34  //     CHECK:   scf.for
35  //     CHECK:     affine.min
36  //     CHECK:     index_cast
37  scf.for %i = %c1 to %c4 step %c2 {
38    %1 = affine.min affine_map<(d0)[s0] -> (2, s0 - d0)> (%i)[%c4]
39    %2 = index_cast %1: index to i64
40    store %2, %A[]: memref<i64>
41  }
42
43  // This should not canonicalize because: 16 - %i may take the value 15 < 1024.
44  //     CHECK:   scf.for
45  //     CHECK:     affine.min
46  //     CHECK:     index_cast
47  scf.for %i = %c1 to %c16 step %c1024 {
48    %1 = affine.min affine_map<(d0) -> (1024, 16 - d0)> (%i)
49    %2 = index_cast %1: index to i64
50    store %2, %A[]: memref<i64>
51  }
52
53  // This example should simplify but affine_map is currently missing
54  // semi-affine canonicalizations: `((s0 * 42 - 1) floordiv s0) * s0`
55  // should evaluate to 41 * s0.
56  // Note that this may require positivity assumptions on `s0`.
57  // Revisit when support is added.
58  // CHECK: scf.for
59  // CHECK:   affine.min
60  // CHECK:   index_cast
61  %ub = affine.apply affine_map<(d0) -> (42 * d0)> (%step)
62  scf.for %i = %c0 to %ub step %step {
63    %1 = affine.min affine_map<(d0, d1, d2) -> (d0, d1 - d2)> (%step, %ub, %i)
64    %2 = index_cast %1: index to i64
65    store %2, %A[]: memref<i64>
66  }
67
68  // This example should simplify but affine_map is currently missing
69  // semi-affine canonicalizations.
70  // This example should simplify but affine_map is currently missing
71  // semi-affine canonicalizations: ` -(((s0 * s0 - 1) floordiv s0) * s0)`
72  // should evaluate to (s0 - 1) * s0.
73  // Note that this may require positivity assumptions on `s0`.
74  // Revisit when support is added.
75  // CHECK: scf.for
76  // CHECK:   affine.min
77  // CHECK:   index_cast
78  %ub2 = affine.apply affine_map<(d0)[s0] -> (s0 * d0)> (%step)[%step]
79  scf.for %i = %c0 to %ub2 step %step {
80    %1 = affine.min affine_map<(d0, d1, d2) -> (d0, d2 - d1)> (%step, %i, %ub2)
81    %2 = index_cast %1: index to i64
82    store %2, %A[]: memref<i64>
83  }
84
85  return
86}
87
88// CHECK-LABEL: scf_parallel
89func @scf_parallel(%A : memref<i64>, %step : index) {
90  %c0 = constant 0 : index
91  %c1 = constant 1 : index
92  %c2 = constant 2 : index
93  %c7 = constant 7 : index
94  %c4 = constant 4 : index
95
96  // CHECK: scf.parallel
97  // CHECK-NEXT:   %[[C2:.*]] = constant 2 : index
98  // CHECK-NEXT:   %[[C2I64:.*]] = index_cast %[[C2:.*]]
99  // CHECK-NEXT:   store %[[C2I64]], %{{.*}}[] : memref<i64>
100  scf.parallel (%i) = (%c0) to (%c4) step (%c2) {
101    %1 = affine.min affine_map<(d0, d1)[] -> (2, d1 - d0)> (%i, %c4)
102    %2 = index_cast %1: index to i64
103    store %2, %A[]: memref<i64>
104  }
105
106  // CHECK: scf.parallel
107  // CHECK-NEXT:   %[[C2:.*]] = constant 2 : index
108  // CHECK-NEXT:   %[[C2I64:.*]] = index_cast %[[C2:.*]]
109  // CHECK-NEXT:   store %[[C2I64]], %{{.*}}[] : memref<i64>
110  scf.parallel (%i) = (%c1) to (%c7) step (%c2) {
111    %1 = affine.min affine_map<(d0)[s0] -> (2, s0 - d0)> (%i)[%c7]
112    %2 = index_cast %1: index to i64
113    store %2, %A[]: memref<i64>
114  }
115
116  // This example should simplify but affine_map is currently missing
117  // semi-affine canonicalizations.
118  // This affine map does not currently evaluate to (0, 0):
119  //   (d0)[s0] -> (s0 mod s0, (-((d0 floordiv s0) * s0) + s0 * 42) mod s0)
120  // TODO: Revisit when support is added.
121  // CHECK: scf.parallel
122  // CHECK:   affine.min
123  // CHECK:   index_cast
124  %ub = affine.apply affine_map<(d0) -> (42 * d0)> (%step)
125  scf.parallel (%i) = (%c0) to (%ub) step (%step) {
126    %1 = affine.min affine_map<(d0, d1, d2) -> (d0, d2 - d1)> (%step, %i, %ub)
127    %2 = index_cast %1: index to i64
128    store %2, %A[]: memref<i64>
129  }
130
131  // This example should simplify but affine_map is currently missing
132  // semi-affine canonicalizations.
133  // This affine map does not currently evaluate to (0, 0):
134  //   (d0)[s0] -> (s0 mod s0, (-((d0 floordiv s0) * s0) + s0 * s0) mod s0)
135  // TODO: Revisit when support is added.
136  // CHECK: scf.parallel
137  // CHECK:   affine.min
138  // CHECK:   index_cast
139  %ub2 = affine.apply affine_map<(d0)[s0] -> (s0 * d0)> (%step)[%step]
140  scf.parallel (%i) = (%c0) to (%ub2) step (%step) {
141    %1 = affine.min affine_map<(d0, d1, d2) -> (d0, d2 - d1)> (%step, %i, %ub2)
142    %2 = index_cast %1: index to i64
143    store %2, %A[]: memref<i64>
144  }
145
146  return
147}
148