1// RUN: mlir-opt -allow-unregistered-dialect -loop-coalescing %s | FileCheck %s 2 3// CHECK-LABEL: @one_3d_nest 4func @one_3d_nest() { 5 // Capture original bounds. Note that for zero-based step-one loops, the 6 // upper bound is also the number of iterations. 7 // CHECK: %[[orig_lb:.*]] = constant 0 8 // CHECK: %[[orig_step:.*]] = constant 1 9 // CHECK: %[[orig_ub_k:.*]] = constant 3 10 // CHECK: %[[orig_ub_i:.*]] = constant 42 11 // CHECK: %[[orig_ub_j:.*]] = constant 56 12 %c0 = constant 0 : index 13 %c1 = constant 1 : index 14 %c2 = constant 2 : index 15 %c3 = constant 3 : index 16 %c42 = constant 42 : index 17 %c56 = constant 56 : index 18 // The range of the new scf. 19 // CHECK: %[[partial_range:.*]] = muli %[[orig_ub_i]], %[[orig_ub_j]] 20 // CHECK-NEXT:%[[range:.*]] = muli %[[partial_range]], %[[orig_ub_k]] 21 22 // Updated loop bounds. 23 // CHECK: scf.for %[[i:.*]] = %[[orig_lb]] to %[[range]] step %[[orig_step]] 24 scf.for %i = %c0 to %c42 step %c1 { 25 // Inner loops must have been removed. 26 // CHECK-NOT: scf.for 27 28 // Reconstruct original IVs from the linearized one. 29 // CHECK: %[[orig_k:.*]] = remi_signed %[[i]], %[[orig_ub_k]] 30 // CHECK: %[[div:.*]] = divi_signed %[[i]], %[[orig_ub_k]] 31 // CHECK: %[[orig_j:.*]] = remi_signed %[[div]], %[[orig_ub_j]] 32 // CHECK: %[[orig_i:.*]] = divi_signed %[[div]], %[[orig_ub_j]] 33 scf.for %j = %c0 to %c56 step %c1 { 34 scf.for %k = %c0 to %c3 step %c1 { 35 // CHECK: "use"(%[[orig_i]], %[[orig_j]], %[[orig_k]]) 36 "use"(%i, %j, %k) : (index, index, index) -> () 37 } 38 } 39 } 40 return 41} 42 43// Check that there is no chasing the replacement of value uses by ensuring 44// multiple uses of loop induction variables get rewritten to the same values. 45 46// CHECK-LABEL: @multi_use 47func @multi_use() { 48 %c0 = constant 0 : index 49 %c1 = constant 1 : index 50 %c10 = constant 10 : index 51 // CHECK: scf.for %[[iv:.*]] = 52 scf.for %i = %c1 to %c10 step %c1 { 53 scf.for %j = %c1 to %c10 step %c1 { 54 scf.for %k = %c1 to %c10 step %c1 { 55 // CHECK: %[[k_unshifted:.*]] = remi_signed %[[iv]], %[[k_extent:.*]] 56 // CHECK: %[[ij:.*]] = divi_signed %[[iv]], %[[k_extent]] 57 // CHECK: %[[j_unshifted:.*]] = remi_signed %[[ij]], %[[j_extent:.*]] 58 // CHECK: %[[i_unshifted:.*]] = divi_signed %[[ij]], %[[j_extent]] 59 // CHECK: %[[k:.*]] = addi %[[k_unshifted]] 60 // CHECK: %[[j:.*]] = addi %[[j_unshifted]] 61 // CHECK: %[[i:.*]] = addi %[[i_unshifted]] 62 63 // CHECK: "use1"(%[[i]], %[[j]], %[[k]]) 64 "use1"(%i,%j,%k) : (index,index,index) -> () 65 // CHECK: "use2"(%[[i]], %[[k]], %[[j]]) 66 "use2"(%i,%k,%j) : (index,index,index) -> () 67 // CHECK: "use3"(%[[k]], %[[j]], %[[i]]) 68 "use3"(%k,%j,%i) : (index,index,index) -> () 69 } 70 } 71 } 72 return 73} 74 75func @unnormalized_loops() { 76 // CHECK: %[[orig_step_i:.*]] = constant 2 77 // CHECK: %[[orig_step_j:.*]] = constant 3 78 // CHECK: %[[orig_lb_i:.*]] = constant 5 79 // CHECK: %[[orig_lb_j:.*]] = constant 7 80 // CHECK: %[[orig_ub_i:.*]] = constant 10 81 // CHECK: %[[orig_ub_j:.*]] = constant 17 82 %c2 = constant 2 : index 83 %c3 = constant 3 : index 84 %c5 = constant 5 : index 85 %c7 = constant 7 : index 86 %c10 = constant 10 : index 87 %c17 = constant 17 : index 88 89 // Number of iterations in the outer scf. 90 // CHECK: %[[diff_i:.*]] = subi %[[orig_ub_i]], %[[orig_lb_i]] 91 // CHECK: %[[c1:.*]] = constant 1 92 // CHECK: %[[step_minus_c1:.*]] = subi %[[orig_step_i]], %[[c1]] 93 // CHECK: %[[dividend:.*]] = addi %[[diff_i]], %[[step_minus_c1]] 94 // CHECK: %[[numiter_i:.*]] = divi_signed %[[dividend]], %[[orig_step_i]] 95 96 // Normalized lower bound and step for the outer scf. 97 // CHECK: %[[lb_i:.*]] = constant 0 98 // CHECK: %[[step_i:.*]] = constant 1 99 100 // Number of iterations in the inner loop, the pattern is the same as above, 101 // only capture the final result. 102 // CHECK: %[[numiter_j:.*]] = divi_signed {{.*}}, %[[orig_step_j]] 103 104 // New bounds of the outer scf. 105 // CHECK: %[[range:.*]] = muli %[[numiter_i]], %[[numiter_j]] 106 // CHECK: scf.for %[[i:.*]] = %[[lb_i]] to %[[range]] step %[[step_i]] 107 scf.for %i = %c5 to %c10 step %c2 { 108 // The inner loop has been removed. 109 // CHECK-NOT: scf.for 110 scf.for %j = %c7 to %c17 step %c3 { 111 // The IVs are rewritten. 112 // CHECK: %[[normalized_j:.*]] = remi_signed %[[i]], %[[numiter_j]] 113 // CHECK: %[[normalized_i:.*]] = divi_signed %[[i]], %[[numiter_j]] 114 // CHECK: %[[scaled_j:.*]] = muli %[[normalized_j]], %[[orig_step_j]] 115 // CHECK: %[[orig_j:.*]] = addi %[[scaled_j]], %[[orig_lb_j]] 116 // CHECK: %[[scaled_i:.*]] = muli %[[normalized_i]], %[[orig_step_i]] 117 // CHECK: %[[orig_i:.*]] = addi %[[scaled_i]], %[[orig_lb_i]] 118 // CHECK: "use"(%[[orig_i]], %[[orig_j]]) 119 "use"(%i, %j) : (index, index) -> () 120 } 121 } 122 return 123} 124 125// Check with parametric loop bounds and steps, capture the bounds here. 126// CHECK-LABEL: @parametric 127// CHECK-SAME: %[[orig_lb1:[A-Za-z0-9]+]]: 128// CHECK-SAME: %[[orig_ub1:[A-Za-z0-9]+]]: 129// CHECK-SAME: %[[orig_step1:[A-Za-z0-9]+]]: 130// CHECK-SAME: %[[orig_lb2:[A-Za-z0-9]+]]: 131// CHECK-SAME: %[[orig_ub2:[A-Za-z0-9]+]]: 132// CHECK-SAME: %[[orig_step2:[A-Za-z0-9]+]]: 133func @parametric(%lb1 : index, %ub1 : index, %step1 : index, 134 %lb2 : index, %ub2 : index, %step2 : index) { 135 // Compute the number of iterations for each of the loops and the total 136 // number of iterations. 137 // CHECK: %[[range1:.*]] = subi %[[orig_ub1]], %[[orig_lb1]] 138 // CHECK: %[[orig_step1_minus_1:.*]] = subi %[[orig_step1]], %c1 139 // CHECK: %[[dividend1:.*]] = addi %[[range1]], %[[orig_step1_minus_1]] 140 // CHECK: %[[numiter1:.*]] = divi_signed %[[dividend1]], %[[orig_step1]] 141 // CHECK: %[[range2:.*]] = subi %[[orig_ub2]], %[[orig_lb2]] 142 // CHECK: %[[orig_step2_minus_1:.*]] = subi %arg5, %c1 143 // CHECK: %[[dividend2:.*]] = addi %[[range2]], %[[orig_step2_minus_1]] 144 // CHECK: %[[numiter2:.*]] = divi_signed %[[dividend2]], %[[orig_step2]] 145 // CHECK: %[[range:.*]] = muli %[[numiter1]], %[[numiter2]] : index 146 147 // Check that the outer loop is updated. 148 // CHECK: scf.for %[[i:.*]] = %c0{{.*}} to %[[range]] step %c1 149 scf.for %i = %lb1 to %ub1 step %step1 { 150 // Check that the inner loop is removed. 151 // CHECK-NOT: scf.for 152 scf.for %j = %lb2 to %ub2 step %step2 { 153 // Remapping of the induction variables. 154 // CHECK: %[[normalized_j:.*]] = remi_signed %[[i]], %[[numiter2]] : index 155 // CHECK: %[[normalized_i:.*]] = divi_signed %[[i]], %[[numiter2]] : index 156 // CHECK: %[[scaled_j:.*]] = muli %[[normalized_j]], %[[orig_step2]] 157 // CHECK: %[[orig_j:.*]] = addi %[[scaled_j]], %[[orig_lb2]] 158 // CHECK: %[[scaled_i:.*]] = muli %[[normalized_i]], %[[orig_step1]] 159 // CHECK: %[[orig_i:.*]] = addi %[[scaled_i]], %[[orig_lb1]] 160 161 // CHECK: "foo"(%[[orig_i]], %[[orig_j]]) 162 "foo"(%i, %j) : (index, index) -> () 163 } 164 } 165 return 166} 167 168// CHECK-LABEL: @two_bands 169func @two_bands() { 170 %c0 = constant 0 : index 171 %c1 = constant 1 : index 172 %c10 = constant 10 : index 173 // CHECK: %[[outer_range:.*]] = muli 174 // CHECK: scf.for %{{.*}} = %{{.*}} to %[[outer_range]] 175 scf.for %i = %c0 to %c10 step %c1 { 176 // Check that the "j" loop was removed and that the inner loops were 177 // coalesced as well. The preparation step for coalescing will inject the 178 // subtraction operation unlike the IV remapping. 179 // CHECK-NOT: scf.for 180 // CHECK: subi 181 scf.for %j = %c0 to %c10 step %c1 { 182 // The inner pair of loops is coalesced separately. 183 // CHECK: scf.for 184 scf.for %k = %i to %j step %c1 { 185 // CHECK_NOT: scf.for 186 scf.for %l = %i to %j step %c1 { 187 "foo"() : () -> () 188 } 189 } 190 } 191 } 192 return 193} 194