mlir/test/Transforms/copy-removal.mlir

   1 // RUN: mlir-opt -copy-removal -split-input-file %s
   2 //| FileCheck %s
   3
   4 // All linalg copies except the linalg.copy(%1, %9) must be removed since the
   5 // defining operation of %1 and its DeallocOp have been defined in another block.
   6
   7 // CHECK-LABEL: func @nested_region_control_flow_div_nested
   8 func @nested_region_control_flow_div_nested(%arg0: index, %arg1: index) -> memref<?x?xf32> {
   9   %0 = cmpi "eq", %arg0, %arg1 : index
  10   %1 = alloc(%arg0, %arg0) : memref<?x?xf32>
  11   // CHECK: %{{.*}} = scf.if
  12   %2 = scf.if %0 -> (memref<?x?xf32>) {
  13     // CHECK: %[[PERCENT3:.*]] = scf.if
  14     %3 = scf.if %0 -> (memref<?x?xf32>) {
  15       %c0_0 = constant 0 : index
  16       %7 = dim %1, %c0_0 : memref<?x?xf32>
  17       %c1_1 = constant 1 : index
  18       %8 = dim %1, %c1_1 : memref<?x?xf32>
  19       %9 = alloc(%7, %8) : memref<?x?xf32>
  20       // CHECK: linalg.copy({{.*}}, %[[PERCENT9:.*]])
  21       linalg.copy(%1, %9) : memref<?x?xf32>, memref<?x?xf32>
  22       // CHECK: scf.yield %[[PERCENT9]]
  23       scf.yield %9 : memref<?x?xf32>
  24     } else {
  25       // CHECK: %[[PERCENT7:.*]] = alloc
  26       %7 = alloc(%arg0, %arg1) : memref<?x?xf32>
  27       %c0_0 = constant 0 : index
  28       %8 = dim %7, %c0_0 : memref<?x?xf32>
  29       %c1_1 = constant 1 : index
  30       %9 = dim %7, %c1_1 : memref<?x?xf32>
  31       // CHECK-NOT: %{{.*}} = alloc
  32       // CHECK-NOT: linalg.copy(%[[PERCENT7]], %{{.*}})
  33       // CHECK-NOT: dealloc %[[PERCENT7]]
  34       %10 = alloc(%8, %9) : memref<?x?xf32>
  35       linalg.copy(%7, %10) : memref<?x?xf32>, memref<?x?xf32>
  36       dealloc %7 : memref<?x?xf32>
  37       // CHECK: scf.yield %[[PERCENT7]]
  38       scf.yield %10 : memref<?x?xf32>
  39     }
  40     %c0 = constant 0 : index
  41     %4 = dim %3, %c0 : memref<?x?xf32>
  42     %c1 = constant 1 : index
  43     %5 = dim %3, %c1 : memref<?x?xf32>
  44     // CHECK-NOT: %{{.*}} = alloc
  45     // CHECK-NOT: linalg.copy(%[[PERCENT3]], %{{.*}})
  46     // CHECK-NOT: dealloc %[[PERCENT3]]
  47     %6 = alloc(%4, %5) : memref<?x?xf32>
  48     linalg.copy(%3, %6) : memref<?x?xf32>, memref<?x?xf32>
  49     dealloc %3 : memref<?x?xf32>
  50     // CHECK: scf.yield %[[PERCENT3]]
  51     scf.yield %6 : memref<?x?xf32>
  52   } else {
  53     // CHECK: %[[PERCENT3:.*]] = alloc
  54     %3 = alloc(%arg1, %arg1) : memref<?x?xf32>
  55     %c0 = constant 0 : index
  56     %4 = dim %3, %c0 : memref<?x?xf32>
  57     %c1 = constant 1 : index
  58     %5 = dim %3, %c1 : memref<?x?xf32>
  59     // CHECK-NOT: %{{.*}} = alloc
  60     // CHECK-NOT: linalg.copy(%[[PERCENT3]], %{{.*}})
  61     // CHECK-NOT: dealloc %[[PERCENT3]]
  62     %6 = alloc(%4, %5) : memref<?x?xf32>
  63     linalg.copy(%3, %6) : memref<?x?xf32>, memref<?x?xf32>
  64     dealloc %3 : memref<?x?xf32>
  65     // CHECK: scf.yield %[[PERCENT3]]
  66     scf.yield %6 : memref<?x?xf32>
  67   }
  68   dealloc %1 : memref<?x?xf32>
  69   return %2 : memref<?x?xf32>
  70 }
  71
  72 // -----
  73
  74 // CHECK-LABEL: func @simple_test
  75 func @simple_test() -> memref<5xf32> {
  76   %temp = alloc() : memref<5xf32>
  77   %ret = alloc() : memref<5xf32>
  78   linalg.copy(%ret, %temp) : memref<5xf32>, memref<5xf32>
  79   dealloc %ret : memref<5xf32>
  80   return %temp : memref<5xf32>
  81 }
  82 // CHECK-SAME: () -> memref<5xf32>
  83 // CHECK-NEXT: %[[ret:.*]] = alloc()
  84 // CHECK-NOT: linalg.copy(%[[ret]], %{{.*}})
  85 // CHECK-NOT: dealloc %[[ret]]
  86 // CHECK: return %[[ret]]
  87
  88 // -----
  89
  90 // It is legal to remove the copy operation that %ret has a usage before the copy
  91 // operation. The allocation of %temp and the deallocation of %ret should be also
  92 // removed.
  93
  94 // CHECK-LABEL: func @test_with_ret_usage_before_copy
  95 func @test_with_ret_usage_before_copy() -> memref<5xf32> {
  96   %ret = alloc() : memref<5xf32>
  97   %temp = alloc() : memref<5xf32>
  98   %c0 = constant 0 : index
  99   %dimension = dim %ret, %c0 : memref<5xf32>
 100   linalg.copy(%ret, %temp) : memref<5xf32>, memref<5xf32>
 101   dealloc %ret : memref<5xf32>
 102   return %temp : memref<5xf32>
 103 }
 104 // CHECK-NEXT: %[[ret:.*]] = alloc()
 105 // CHECK-NOT: %{{.*}} = alloc
 106 // CHECK-NEXT: %{{.*}} = constant
 107 // CHECK-NEXT: %[[DIM:.*]] = dim %[[ret]]
 108 // CHECK-NOT: linalg.copy(%[[ret]], %{{.*}})
 109 // CHECK-NOT: dealloc %[[ret]]
 110 // CHECK: return %[[ret]]
 111
 112 // -----
 113
 114 // It is illegal to remove a copy operation that %ret has a usage after copy
 115 // operation.
 116
 117 // CHECK-LABEL: func @test_with_ret_usage_after_copy
 118 func @test_with_ret_usage_after_copy() -> memref<5xf32> {
 119   %ret = alloc() : memref<5xf32>
 120   %temp = alloc() : memref<5xf32>
 121   // CHECK: linalg.copy
 122   linalg.copy(%ret, %temp) : memref<5xf32>, memref<5xf32>
 123   %c0 = constant 0 : index
 124   %dimension = dim %ret, %c0 : memref<5xf32>
 125   dealloc %ret : memref<5xf32>
 126   return %temp : memref<5xf32>
 127 }
 128
 129 // -----
 130
 131 // It is illegal to remove a copy operation that %temp has a usage before copy
 132 // operation.
 133
 134 // CHECK-LABEL: func @test_with_temp_usage_before_copy
 135 func @test_with_temp_usage_before_copy() -> memref<5xf32> {
 136   %ret = alloc() : memref<5xf32>
 137   %temp = alloc() : memref<5xf32>
 138   %c0 = constant 0 : index
 139   %dimension = dim %temp, %c0 : memref<5xf32>
 140   // CHECK: linalg.copy
 141   linalg.copy(%ret, %temp) : memref<5xf32>, memref<5xf32>
 142   dealloc %ret : memref<5xf32>
 143   return %temp : memref<5xf32>
 144 }
 145
 146 // -----
 147
 148 // It is legal to remove the copy operation that %temp has a usage after the copy
 149 // operation. The allocation of %temp and the deallocation of %ret could be also
 150 // removed.
 151
 152 // However the following pattern is not handled by copy removal.
 153 //   %from = alloc()
 154 //   %to = alloc()
 155 //   copy(%from, %to)
 156 //   read_from(%from) + write_to(%something_else)
 157 //   dealloc(%from)
 158 //   return %to
 159 // In particular, linalg.generic is a memoryEffectOp between copy and dealloc.
 160 // Since no alias analysis is performed and no distinction is made between reads
 161 // and writes, the linalg.generic with effects blocks copy removal.
 162
 163 #map0 = affine_map<(d0) -> (d0)>
 164
 165 // CHECK-LABEL: func @test_with_temp_usage_after_copy
 166 func @test_with_temp_usage_after_copy() -> memref<5xf32> {
 167   %ret = alloc() : memref<5xf32>
 168   %res = alloc() : memref<5xf32>
 169   %temp = alloc() : memref<5xf32>
 170   linalg.copy(%ret, %temp) : memref<5xf32>, memref<5xf32>
 171   linalg.generic {
 172     indexing_maps = [#map0, #map0],
 173     iterator_types = ["parallel"]}
 174     ins(%temp : memref<5xf32>)
 175    outs(%res : memref<5xf32>) {
 176   ^bb0(%gen1_arg0: f32, %gen1_arg1: f32):
 177     %tmp1 = exp %gen1_arg0 : f32
 178     linalg.yield %tmp1 : f32
 179   }
 180   dealloc %ret : memref<5xf32>
 181   return %temp : memref<5xf32>
 182 }
 183 // CHECK-NEXT: %[[ret:.*]] = alloc()
 184 // CHECK-NEXT: %[[res:.*]] = alloc()
 185 // CHECK-NEXT: %[[temp:.*]] = alloc()
 186 // CHECK-NEXT: linalg.copy(%[[ret]], %[[temp]])
 187 // CHECK-NEXT: linalg.generic
 188 //      CHECK: dealloc %[[ret]]
 189 //      CHECK: return %[[temp]]
 190
 191 // -----
 192
 193 // CHECK-LABEL: func @make_allocation
 194 func @make_allocation() -> memref<5xf32> {
 195   %mem = alloc() : memref<5xf32>
 196   return %mem : memref<5xf32>
 197 }
 198
 199 // CHECK-LABEL: func @test_with_function_call
 200 func @test_with_function_call() -> memref<5xf32> {
 201   // CHECK-NEXT: %[[ret:.*]] = call @make_allocation() : () -> memref<5xf32>
 202   %ret = call @make_allocation() : () -> (memref<5xf32>)
 203   // CHECK-NOT: %{{.*}} = alloc
 204   // CHECK-NOT: linalg.copy(%[[ret]], %{{.*}})
 205   // CHECK-NOT: dealloc %[[ret]]
 206   %temp = alloc() : memref<5xf32>
 207   linalg.copy(%ret, %temp) : memref<5xf32>, memref<5xf32>
 208   dealloc %ret : memref<5xf32>
 209   // CHECK: return %[[ret]]
 210   return %temp : memref<5xf32>
 211 }
 212
 213 // -----
 214
 215 // CHECK-LABEL: func @multiple_deallocs_in_different_blocks
 216 func @multiple_deallocs_in_different_blocks(%cond : i1) -> memref<5xf32> {
 217   // CHECK-NEXT: %[[PERCENT0:.*]] = alloc()
 218   %0 = alloc() : memref<5xf32>
 219   cond_br %cond, ^bb1, ^bb2
 220 ^bb1:
 221   dealloc %0 : memref<5xf32>
 222   // CHECK: br ^[[BB3:.*]](%[[PERCENT0]]
 223   br ^bb3(%0 : memref<5xf32>)
 224 ^bb2:
 225   // CHECK-NOT: %{{.*}} = alloc
 226   // CHECK-NOT: linalg.copy(%[[PERCENT0]], %{{.*}})
 227   // CHECK-NOT: dealloc %[[PERCENT0]]
 228   %temp = alloc() : memref<5xf32>
 229   linalg.copy(%0, %temp) : memref<5xf32>, memref<5xf32>
 230   dealloc %0 : memref<5xf32>
 231   // CHECK: br ^[[BB3]](%[[PERCENT0]]
 232   br ^bb3(%temp : memref<5xf32>)
 233 ^bb3(%res : memref<5xf32>):
 234   return %res : memref<5xf32>
 235 }
 236
 237 // -----
 238
 239 #map0 = affine_map<(d0) -> (d0)>
 240
 241 // CHECK-LABEL: func @test_ReuseCopyTargetAsSource
 242 func @test_ReuseCopyTargetAsSource(%arg0: memref<2xf32>, %result: memref<2xf32>){
 243   // CHECK-SAME: (%[[ARG0:.*]]: memref<2xf32>, %[[RES:.*]]: memref<2xf32>)
 244   // CHECK-NOT: %{{.*}} = alloc
 245   %temp = alloc() : memref<2xf32>
 246   // CHECK-NEXT: linalg.generic
 247   // CHECK-SAME: ins(%[[ARG0]]{{.*}}outs(%[[RES]]
 248   // CHECK-NOT: linalg.copy(%{{.*}}, %[[RES]])
 249   // CHECK-NOT: dealloc %{{.*}}
 250   linalg.generic {
 251     indexing_maps = [#map0, #map0],
 252     iterator_types = ["parallel"]}
 253     ins(%arg0 : memref<2xf32>)
 254    outs(%temp : memref<2xf32>) {
 255   ^bb0(%gen2_arg0: f32, %gen2_arg1: f32):
 256     %tmp2 = exp %gen2_arg0 : f32
 257     linalg.yield %tmp2 : f32
 258   }
 259   "linalg.copy"(%temp, %result) : (memref<2xf32>, memref<2xf32>) -> ()
 260   dealloc %temp : memref<2xf32>
 261   // CHECK: return
 262   return
 263 }
 264
 265 // -----
 266
 267 // Copy operation must not be removed since an operation writes to %to value
 268 // before copy.
 269
 270 #map0 = affine_map<(d0) -> (d0)>
 271
 272 // CHECK-LABEL: func @test_ReuseCopyTargetAsSource
 273 func @test_ReuseCopyTargetAsSource(%arg0: memref<2xf32>){
 274   %to = alloc() : memref<2xf32>
 275   %temp = alloc() : memref<2xf32>
 276   linalg.generic {
 277     indexing_maps = [#map0, #map0],
 278     iterator_types = ["parallel"]}
 279     ins(%arg0 : memref<2xf32>)
 280    outs(%temp : memref<2xf32>) {
 281   ^bb0(%gen1_arg0: f32, %gen1_arg1: f32):
 282     %tmp1 = exp %gen1_arg0 : f32
 283     linalg.yield %tmp1 : f32
 284   }
 285   linalg.generic {
 286     indexing_maps = [#map0, #map0],
 287     iterator_types = ["parallel"]}
 288     ins(%arg0 : memref<2xf32>)
 289    outs(%to : memref<2xf32>) {
 290   ^bb0(%gen2_arg0: f32, %gen2_arg1: f32):
 291     %tmp2 = exp %gen2_arg0 : f32
 292     linalg.yield %tmp2 : f32
 293   }
 294   // CHECK: linalg.copy
 295   "linalg.copy"(%temp, %to) : (memref<2xf32>, memref<2xf32>) -> ()
 296   dealloc %temp : memref<2xf32>
 297   return
 298 }
 299
 300 // -----
 301
 302 // The only redundant copy is linalg.copy(%4, %5)
 303
 304 // CHECK-LABEL: func @loop_alloc
 305 func @loop_alloc(%arg0: index, %arg1: index, %arg2: index, %arg3: memref<2xf32>, %arg4: memref<2xf32>) {
 306   // CHECK: %{{.*}} = alloc()
 307   %0 = alloc() : memref<2xf32>
 308   dealloc %0 : memref<2xf32>
 309   // CHECK: %{{.*}} = alloc()
 310   %1 = alloc() : memref<2xf32>
 311   // CHECK: linalg.copy
 312   linalg.copy(%arg3, %1) : memref<2xf32>, memref<2xf32>
 313   %2 = scf.for %arg5 = %arg0 to %arg1 step %arg2 iter_args(%arg6 = %1) -> (memref<2xf32>) {
 314     %3 = cmpi "eq", %arg5, %arg1 : index
 315     // CHECK: dealloc
 316     dealloc %arg6 : memref<2xf32>
 317     // CHECK: %[[PERCENT4:.*]] = alloc()
 318     %4 = alloc() : memref<2xf32>
 319     // CHECK-NOT: alloc
 320     // CHECK-NOT: linalg.copy
 321     // CHECK-NOT: dealloc
 322     %5 = alloc() : memref<2xf32>
 323     linalg.copy(%4, %5) : memref<2xf32>, memref<2xf32>
 324     dealloc %4 : memref<2xf32>
 325     // CHECK: %[[PERCENT6:.*]] = alloc()
 326     %6 = alloc() : memref<2xf32>
 327     // CHECK: linalg.copy(%[[PERCENT4]], %[[PERCENT6]])
 328     linalg.copy(%5, %6) : memref<2xf32>, memref<2xf32>
 329     scf.yield %6 : memref<2xf32>
 330   }
 331   // CHECK: linalg.copy
 332   linalg.copy(%2, %arg4) : memref<2xf32>, memref<2xf32>
 333   dealloc %2 : memref<2xf32>
 334   return
 335 }
 336
 337 // -----
 338
 339 // The linalg.copy operation can be removed in addition to alloc and dealloc
 340 // operations. All uses of %0 is then replaced with %arg2.
 341
 342 // CHECK-LABEL: func @check_with_affine_dialect
 343 func @check_with_affine_dialect(%arg0: memref<4xf32>, %arg1: memref<4xf32>, %arg2: memref<4xf32>) {
 344   // CHECK-SAME: (%[[ARG0:.*]]: memref<4xf32>, %[[ARG1:.*]]: memref<4xf32>, %[[RES:.*]]: memref<4xf32>)
 345   // CHECK-NOT: alloc
 346   %0 = alloc() : memref<4xf32>
 347   affine.for %arg3 = 0 to 4 {
 348     %5 = affine.load %arg0[%arg3] : memref<4xf32>
 349     %6 = affine.load %arg1[%arg3] : memref<4xf32>
 350     %7 = cmpf "ogt", %5, %6 : f32
 351     // CHECK: %[[SELECT_RES:.*]] = select
 352     %8 = select %7, %5, %6 : f32
 353     // CHECK-NEXT: affine.store %[[SELECT_RES]], %[[RES]]
 354     affine.store %8, %0[%arg3] : memref<4xf32>
 355   }
 356   // CHECK-NOT: linalg.copy
 357   // CHECK-NOT: dealloc
 358   "linalg.copy"(%0, %arg2) : (memref<4xf32>, memref<4xf32>) -> ()
 359   dealloc %0 : memref<4xf32>
 360   //CHECK: return
 361   return
 362 }