mlir/test/Dialect/Affine/loop-unswitch.mlir

   1 // RUN: mlir-opt %s -split-input-file -pass-pipeline="builtin.module(func.func(test-affine-loop-unswitch))" | FileCheck %s
   2
   3 // CHECK-DAG: #[[$SET:.*]] = affine_set<(d0) : (d0 - 2 >= 0)>
   4
   5 // CHECK-LABEL: func @if_else_imperfect
   6 func.func @if_else_imperfect(%A : memref<100xi32>, %B : memref<100xi32>, %v : i32) {
   7 // CHECK: %[[A:.*]]: memref<100xi32>, %[[B:.*]]: memref
   8   affine.for %i = 0 to 100 {
   9     affine.store %v, %A[%i] : memref<100xi32>
  10     affine.for %j = 0 to 100 {
  11       affine.store %v, %A[%j] : memref<100xi32>
  12       affine.if affine_set<(d0) : (d0 - 2 >= 0)>(%i) {
  13         affine.store %v, %B[%j] : memref<100xi32>
  14       }
  15       func.call @external() : () -> ()
  16     }
  17     affine.store %v, %A[%i] : memref<100xi32>
  18   }
  19   return
  20 }
  21 func.func private @external()
  22
  23 // CHECK:       affine.for %[[I:.*]] = 0 to 100 {
  24 // CHECK-NEXT:    affine.store %{{.*}}, %[[A]][%[[I]]]
  25 // CHECK-NEXT:    affine.if #[[$SET]](%[[I]]) {
  26 // CHECK-NEXT:      affine.for %[[J:.*]] = 0 to 100 {
  27 // CHECK-NEXT:        affine.store %{{.*}}, %[[A]][%[[J]]]
  28 // CHECK-NEXT:        affine.store %{{.*}}, %[[B]][%[[J]]]
  29 // CHECK-NEXT:        call
  30 // CHECK-NEXT:      }
  31 // CHECK-NEXT:    } else {
  32 // CHECK-NEXT:      affine.for %[[JJ:.*]] = 0 to 100 {
  33 // CHECK-NEXT:        affine.store %{{.*}}, %[[A]][%[[J]]]
  34 // CHECK-NEXT:        call
  35 // CHECK-NEXT:      }
  36 // CHECK-NEXT:    }
  37 // CHECK-NEXT:    affine.store %{{.*}}, %[[A]][%[[I]]]
  38 // CHECK-NEXT:  }
  39 // CHECK-NEXT:  return
  40
  41 // -----
  42
  43 func.func private @foo()
  44 func.func private @bar()
  45 func.func private @abc()
  46 func.func private @xyz()
  47
  48 // CHECK-LABEL: func @if_then_perfect
  49 func.func @if_then_perfect(%A : memref<100xi32>, %v : i32) {
  50   affine.for %i = 0 to 100 {
  51     affine.for %j = 0 to 100 {
  52       affine.for %k = 0 to 100 {
  53         affine.if affine_set<(d0) : (d0 - 2 >= 0)>(%i) {
  54           affine.store %v, %A[%i] : memref<100xi32>
  55         }
  56       }
  57     }
  58   }
  59   return
  60 }
  61 // CHECK:      affine.for
  62 // CHECK-NEXT:   affine.if
  63 // CHECK-NEXT:     affine.for
  64 // CHECK-NEXT:       affine.for
  65 // CHECK-NOT:    else
  66
  67
  68 // CHECK-LABEL: func @if_else_perfect
  69 func.func @if_else_perfect(%A : memref<100xi32>, %v : i32) {
  70   affine.for %i = 0 to 99 {
  71     affine.for %j = 0 to 100 {
  72       affine.for %k = 0 to 100 {
  73         func.call @foo() : () -> ()
  74         affine.if affine_set<(d0, d1) : (d0 - 2 >= 0, -d1 + 80 >= 0)>(%i, %j) {
  75           affine.store %v, %A[%i] : memref<100xi32>
  76           func.call @abc() : () -> ()
  77         } else {
  78           affine.store %v, %A[%i + 1] : memref<100xi32>
  79           func.call @xyz() : () -> ()
  80         }
  81         func.call @bar() : () -> ()
  82       }
  83     }
  84   }
  85   return
  86 }
  87 // CHECK:      affine.for
  88 // CHECK-NEXT:   affine.for
  89 // CHECK-NEXT:     affine.if
  90 // CHECK-NEXT:       affine.for
  91 // CHECK-NEXT:         call @foo
  92 // CHECK-NEXT:         affine.store %{{.*}}, %{{.*}}[%{{.*}}]
  93 // CHECK-NEXT:         call @abc
  94 // CHECK-NEXT:         call @bar
  95 // CHECK-NEXT:       }
  96 // CHECK-NEXT:     else
  97 // CHECK-NEXT:       affine.for
  98 // CHECK-NEXT:         call @foo
  99 // CHECK-NEXT:         affine.store %{{.*}}, %{{.*}}[%{{.*}} + 1]
 100 // CHECK-NEXT:         call @xyz
 101 // CHECK-NEXT:         call @bar
 102 // CHECK-NEXT:       }
 103 // CHECK-NEXT:     }
 104 // CHECK-NEXT:   }
 105 // CHECK-NEXT: }
 106
 107 // CHECK-LABEL: func @if_then_imperfect
 108 func.func @if_then_imperfect(%A : memref<100xi32>, %N : index, %v: i32) {
 109   affine.for %i = 0 to 100 {
 110     affine.store %v, %A[0] : memref<100xi32>
 111     affine.if affine_set<(d0) : (d0 - 2 >= 0)>(%N) {
 112       affine.store %v, %A[%i] : memref<100xi32>
 113     }
 114   }
 115   return
 116 }
 117 // CHECK:       affine.if
 118 // CHECK-NEXT:    affine.for
 119 // CHECK-NEXT:      affine.store
 120 // CHECK-NEXT:      affine.store
 121 // CHECK-NEXT:    }
 122 // CHECK-NEXT:  } else {
 123 // CHECK-NEXT:    affine.for
 124 // CHECK-NEXT:      affine.store
 125 // CHECK-NEXT:    }
 126 // CHECK-NEXT:  }
 127 // CHECK-NEXT:  return
 128
 129 // Check if unused operands are dropped: hence, hoisting is possible.
 130 // CHECK-LABEL: func @hoist_after_canonicalize
 131 func.func @hoist_after_canonicalize() {
 132   affine.for %i = 0 to 100 {
 133     affine.for %j = 0 to 100 {
 134       affine.if affine_set<(d0) : (d0 - 2 >= 0)>(%j)  {
 135         affine.if affine_set<(d0, d1) : (d0 - 1 >= 0, -d0 + 99 >= 0)>(%i, %j)  {
 136           // The call to external is to avoid DCE on affine.if.
 137           func.call @foo() : () -> ()
 138         }
 139       }
 140     }
 141   }
 142   return
 143 }
 144 // CHECK:      affine.for
 145 // CHECK-NEXT:   affine.if
 146 // CHECK-NEXT:     affine.for
 147 // CHECK-NEXT:       affine.if
 148 // CHECK-NEXT:         call
 149 // CHECK-NEXT:       }
 150 // CHECK-NEXT:     }
 151 // CHECK-NEXT:   }
 152 // CHECK-NEXT: }
 153 // CHECK-NEXT: return
 154
 155 // CHECK-LABEL: func @handle_dead_if
 156 func.func @handle_dead_if(%N : index) {
 157   affine.for %i = 0 to 100 {
 158       affine.if affine_set<(d0) : (d0 - 1 >= 0, -d0 + 99 >= 0)>(%N)  {
 159       }
 160   }
 161   return
 162 }
 163 // CHECK-NEXT: affine.for
 164 // CHECK-NEXT: }
 165 // CHECK-NEXT: return
 166
 167 // -----
 168
 169 // A test case with affine.parallel.
 170
 171 #flb1 = affine_map<(d0) -> (d0 * 3)>
 172 #fub1 = affine_map<(d0) -> (d0 * 3 + 3)>
 173 #flb0 = affine_map<(d0) -> (d0 * 16)>
 174 #fub0 = affine_map<(d0) -> (d0 * 16 + 16)>
 175 #pub1 = affine_map<(d0)[s0] -> (s0, d0 * 3 + 3)>
 176 #pub0 = affine_map<(d0)[s0] -> (s0, d0 * 16 + 16)>
 177 #lb1 = affine_map<(d0) -> (d0 * 480)>
 178 #ub1 = affine_map<(d0)[s0] -> (s0, d0 * 480 + 480)>
 179 #lb0 = affine_map<(d0) -> (d0 * 110)>
 180 #ub0 = affine_map<(d0)[s0] -> (d0 * 110 + 110, s0 floordiv 3)>
 181
 182 #set0 = affine_set<(d0, d1)[s0, s1] : (d0 * -16 + s0 - 16 >= 0, d1 * -3 + s1 - 3 >= 0)>
 183
 184 // CHECK-LABEL: func @perfect_if_else
 185 func.func @perfect_if_else(%arg0 : memref<?x?xf64>, %arg1 : memref<?x?xf64>, %v : f64,
 186             %arg4 : index, %arg5 : index, %arg6 : index, %sym : index) {
 187   affine.for %arg7 = #lb0(%arg5) to min #ub0(%arg5)[%sym] {
 188     affine.parallel (%i0, %j0) = (0, 0) to (symbol(%sym), 100) step (10, 10) {
 189       affine.for %arg8 = #lb1(%arg4) to min #ub1(%arg4)[%sym] {
 190         affine.if #set0(%arg6, %arg7)[%sym, %sym] {
 191           affine.for %arg9 = #flb0(%arg6) to #fub0(%arg6) {
 192             affine.for %arg10 = #flb1(%arg7) to #fub1(%arg7) {
 193               affine.store %v, %arg0[0, 0] : memref<?x?xf64>
 194             }
 195           }
 196         } else {
 197           affine.for %arg9 = #lb0(%arg6) to min #pub0(%arg6)[%sym] {
 198             affine.for %arg10 = #lb1(%arg7) to min #pub1(%arg7)[%sym] {
 199               affine.store %v, %arg0[0, 0] : memref<?x?xf64>
 200             }
 201           }
 202         }
 203       }
 204     }
 205   }
 206   return
 207 }
 208
 209 // CHECK:       affine.for
 210 // CHECK-NEXT:    affine.if
 211 // CHECK-NEXT:      affine.parallel
 212 // CHECK-NEXT:        affine.for
 213 // CHECK-NEXT:          affine.for
 214 // CHECK-NEXT:            affine.for
 215 // CHECK-NEXT:              affine.store
 216 // CHECK-NEXT:            }
 217 // CHECK-NEXT:          }
 218 // CHECK-NEXT:        }
 219 // CHECK-NEXT:      }
 220 // CHECK-NEXT:    } else {
 221 // CHECK-NEXT:      affine.parallel
 222 // CHECK-NEXT:        affine.for
 223 // CHECK-NEXT:          affine.for
 224 // CHECK-NEXT:            affine.for
 225 // CHECK-NEXT:              affine.store
 226 // CHECK-NEXT:            }
 227 // CHECK-NEXT:          }
 228 // CHECK-NEXT:        }
 229 // CHECK-NEXT:      }
 230 // CHECK-NEXT:    }
 231 // CHECK-NEXT:  }
 232
 233 // With multiple if ops in a function, the test pass just looks for the first if
 234 // op that it is able to successfully hoist.
 235
 236 // CHECK-LABEL: func @multiple_if
 237 func.func @multiple_if(%N : index) {
 238   affine.if affine_set<() : (0 == 0)>() {
 239     func.call @external() : () -> ()
 240   }
 241   affine.for %i = 0 to 100 {
 242     affine.if affine_set<()[s0] : (s0 >= 0)>()[%N] {
 243       func.call @external() : () -> ()
 244     }
 245   }
 246   return
 247 }
 248 // CHECK:      call
 249 // CHECK-NEXT: affine.if
 250 // CHECK-NEXT:   affine.for
 251 // CHECK-NEXT:     call
 252 // CHECK-NEXT:   }
 253 // CHECK-NEXT: }
 254 // CHECK-NEXT: return
 255
 256 func.func private @external()