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