mlir/test/Transforms/loop-coalescing.mlir

   1 // RUN: mlir-opt -allow-unregistered-dialect -loop-coalescing %s | FileCheck %s
   2
   3 // CHECK-LABEL: @one_3d_nest
   4 func @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
  47 func @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
  75 func @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]+]]:
 133 func @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
 169 func @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 }