mlir/test/Transforms/parametric-tiling.mlir

   1 // RUN: mlir-opt -test-extract-fixed-outer-loops='test-outer-loop-sizes=7' %s | FileCheck %s --check-prefixes=COMMON,TILE_7
   2 // RUN: mlir-opt -test-extract-fixed-outer-loops='test-outer-loop-sizes=7,4' %s | FileCheck %s --check-prefixes=COMMON,TILE_74
   3
   4 // COMMON-LABEL: @rectangular
   5 func.func @rectangular(%arg0: memref<?x?xf32>) {
   6   %c2 = arith.constant 2 : index
   7   %c44 = arith.constant 44 : index
   8   %c1 = arith.constant 1 : index
   9   // Range of the original loop:
  10   //   (upper - lower + step - 1) / step
  11   // where step is known to be %c1.
  12   // COMMON:      %[[diff:.*]] = arith.subi %c44, %c2
  13   // COMMON:      %[[adjustment:.*]] = arith.subi %c1, %c1_{{.*}}
  14   // COMMON-NEXT: %[[diff_adj:.*]] = arith.addi %[[diff]], %[[adjustment]]
  15   // COMMON-NEXT: %[[range:.*]] = arith.divui %[[diff_adj]], %c1
  16
  17   // Ceildiv to get the parametric tile size.
  18   // COMMON:       %[[sum:.*]] = arith.addi %[[range]], %c6
  19   // COMMON-NEXT:  %[[size:.*]] = arith.divui %[[sum]], %c7
  20   // New outer step (original is %c1).
  21   // COMMON-NEXT:      %[[step:.*]] = arith.muli %c1, %[[size]]
  22
  23   // Range of the second original loop
  24   //   (upper - lower + step - 1) / step
  25   // where step is known to be %c2.
  26   // TILE_74:      %[[diff2:.*]] = arith.subi %c44, %c1
  27   // TILE_74:      %[[adjustment2:.*]] = arith.subi %c2, %c1_{{.*}}
  28   // TILE_74-NEXT: %[[diff2_adj:.*]] = arith.addi %[[diff2]], %[[adjustment2]]
  29   // TILE_74-NEXT: %[[range2:.*]] = arith.divui %[[diff2_adj]], %c2
  30
  31   // Ceildiv to get the parametric tile size for the second original scf.
  32   // TILE_74:      %[[sum2:.*]] = arith.addi %[[range2]], %c3
  33   // TILE_74-NEXT: %[[size2:.*]] = arith.divui %[[sum2]], %c4
  34   // New inner step (original is %c2).
  35   // TILE_74-NEXT:     %[[step2:.*]] = arith.muli %c2, %[[size2]]
  36
  37   // Updated outer loop(s) use new steps.
  38   // COMMON: scf.for %[[i:.*]] = %c2 to %c44 step %[[step]]
  39   // TILE_74:scf.for %[[j:.*]] = %c1 to %c44 step %[[step2]]
  40  scf.for %i = %c2 to %c44 step %c1 {
  41     // Upper bound for the inner loop min(%i + %step, %c44).
  42     // COMMON:      %[[stepped:.*]] = arith.addi %[[i]], %[[step]]
  43     // COMMON-NEXT: %[[ub:.*]] = arith.minsi %c44, %[[stepped]]
  44     //
  45     // TILE_74:      %[[stepped2:.*]] = arith.addi %[[j]], %[[step2]]
  46     // TILE_74-NEXT: %[[ub2:.*]] = arith.minsi %c44, %[[stepped2]]
  47
  48     // Created inner scf.
  49     // COMMON:scf.for %[[ii:.*]] = %[[i]] to %[[ub:.*]] step %c1
  50
  51     // This loop is not modified in TILE_7 case.
  52     // TILE_7: scf.for %[[j:.*]] = %c1 to %c44 step %c2
  53     //
  54     // But is modified in TILE_74 case.
  55     // TILE_74:scf.for %[[jj:.*]] = %[[j]] to %[[ub2]] step %c2
  56    scf.for %j = %c1 to %c44 step %c2 {
  57       // The right iterator are used.
  58       // TILE_7:  memref.load %arg0[%[[ii]], %[[j]]]
  59       // TILE_74: memref.load %arg0[%[[ii]], %[[jj]]]
  60       memref.load %arg0[%i, %j]: memref<?x?xf32>
  61     }
  62   }
  63   return
  64 }
  65
  66 // COMMON-LABEL: @triangular
  67 func.func @triangular(%arg0: memref<?x?xf32>) {
  68   %c2 = arith.constant 2 : index
  69   %c44 = arith.constant 44 : index
  70   %c1 = arith.constant 1 : index
  71   // Range of the original outer loop:
  72   //   (upper - lower + step - 1) / step
  73   // where step is known to be %c1.
  74   // COMMON:      %[[diff:.*]] = arith.subi %c44, %c2
  75   // COMMON:      %[[adjustment:.*]] = arith.subi %c1, %c1_{{.*}}
  76   // COMMON-NEXT: %[[diff_adj:.*]] = arith.addi %[[diff]], %[[adjustment]]
  77   // COMMON-NEXT: %[[range:.*]] = arith.divui %[[diff_adj]], %c1
  78
  79   // Ceildiv to get the parametric tile size.
  80   // COMMON:       %[[sum:.*]] = arith.addi %[[range]], %c6
  81   // COMMON-NEXT:  %[[size:.*]] = arith.divui %[[sum]], %c7
  82   // New outer step (original is %c1).
  83   // COMMON-NEXT:  %[[step:.*]] = arith.muli %c1, %[[size]]
  84
  85   // Constant adjustment for inner loop has been hoisted out.
  86   // TILE_74:      %[[adjustment2:.*]] = arith.subi %c2, %c1_{{.*}}
  87
  88   // New outer scf.
  89   // COMMON: scf.for %[[i:.*]] = %c2 to %c44 step %[[step]]
  90
  91   // Range of the original inner loop
  92   //   (upper - lower + step - 1) / step
  93   // where step is known to be %c2.
  94   // TILE_74:      %[[diff2:.*]] = arith.subi %[[i]], %c1
  95   // TILE_74-NEXT: %[[diff2_adj:.*]] = arith.addi %[[diff2]], %[[adjustment2]]
  96   // TILE_74-NEXT: %[[range2:.*]] = arith.divui %[[diff2_adj]], %c2
  97
  98   // Ceildiv to get the parametric tile size for the second original scf.
  99   // TILE_74:      %[[sum2:.*]] = arith.addi %[[range2]], %c3
 100   // TILE_74-NEXT: %[[size2:.*]] = arith.divui %[[sum2]], %c4
 101   // New inner step (original is %c2).
 102   // TILE_74-NEXT:     %[[step2:.*]] = arith.muli %c2, %[[size2]]
 103
 104   // New inner scf.
 105   // TILE_74:scf.for %[[j:.*]] = %c1 to %[[i]] step %[[step2]]
 106  scf.for %i = %c2 to %c44 step %c1 {
 107     // Upper bound for the inner loop min(%i + %step, %c44).
 108     // COMMON:      %[[stepped:.*]] = arith.addi %[[i]], %[[step]]
 109     // COMMON-NEXT: %[[ub:.*]] = arith.minsi %c44, %[[stepped]]
 110     // TILE_74:      %[[stepped2:.*]] = arith.addi %[[j]], %[[step2]]
 111     // TILE_74-NEXT: %[[ub2:.*]] = arith.minsi %[[i]], %[[stepped2]]
 112     //
 113     // Created inner scf.
 114     // COMMON:scf.for %[[ii:.*]] = %[[i]] to %[[ub:.*]] step %c1
 115
 116     // This loop is not modified in TILE_7 case.
 117     // TILE_7: scf.for %[[j:.*]] = %c1 to %[[ii]] step %c2
 118     //
 119     // But is modified in TILE_74 case.
 120     // TILE_74:scf.for %[[jj:.*]] = %[[j]] to %[[ub2]] step %c2
 121    scf.for %j = %c1 to %i step %c2 {
 122       // The right iterator are used.
 123       // TILE_7:  memref.load %arg0[%[[ii]], %[[j]]]
 124       // TILE_74: memref.load %arg0[%[[ii]], %[[jj]]]
 125       memref.load %arg0[%i, %j]: memref<?x?xf32>
 126     }
 127   }
 128   return
 129 }