178 lines · plain
1// RUN: mlir-opt %s -split-input-file -allow-unregistered-dialect -pass-pipeline="builtin.module(func.func(linalg-detensorize))" | FileCheck %s2 3#map0 = affine_map<() -> ()>4 5#attrs = {6 indexing_maps = [#map0, #map0, #map0],7 iterator_types = []8}9 10func.func @main() -> (tensor<i32>) attributes {} {11 %c0 = arith.constant 0 : i3212 %0 = tensor.from_elements %c0 : tensor<i32>13 %c10 = arith.constant 10 : i3214 %1 = tensor.from_elements %c10 : tensor<i32>15 cf.br ^bb1(%0 : tensor<i32>)16 17^bb1(%2: tensor<i32>): // 2 preds: ^bb0, ^bb218 %3 = tensor.empty() : tensor<i1>19 %4 = linalg.generic #attrs20 ins(%2, %1 : tensor<i32>, tensor<i32>)21 outs(%3 : tensor<i1>) {22 ^bb0(%arg0: i32, %arg1: i32, %arg2: i1):23 %8 = arith.cmpi slt, %arg0, %arg1 : i3224 linalg.yield %8 : i125 } -> tensor<i1>26 %5 = tensor.extract %4[] : tensor<i1>27 cf.cond_br %5, ^bb2(%2 : tensor<i32>), ^bb3(%2 : tensor<i32>)28 29^bb2(%6: tensor<i32>): // pred: ^bb130 %7 = tensor.empty() : tensor<i32>31 %8 = linalg.generic #attrs32 ins(%6, %6 : tensor<i32>, tensor<i32>)33 outs(%7 : tensor<i32>) {34 ^bb0(%arg0: i32, %arg1: i32, %arg2: i32):35 %9 = arith.addi %arg0, %arg1 : i3236 linalg.yield %9 : i3237 } -> tensor<i32>38 cf.br ^bb3(%8 : tensor<i32>)39 40^bb3(%10: tensor<i32>): // pred: ^bb141 return %10 : tensor<i32>42}43 44// CHECK-LABEL: func @main()45// CHECK-DAG: %[[cst:.*]] = arith.constant dense<0>46// CHECK-DAG: arith.constant true47// CHECK: cf.br48// CHECK-NEXT: ^[[bb1:.*]]:49// CHECK-NEXT: cf.cond_br %{{.*}}, ^[[bb2:.*]], ^bb350// CHECK-NEXT: ^[[bb2]]51// CHECK-NEXT: cf.br ^[[bb3:.*]]52// CHECK-NEXT: ^[[bb3]]53// CHECK-NEXT: return %[[cst]]54// CHECK-NEXT: }55 56// -----57 58// Similar to the above test with one change: one of the block after the59// if-condition passes/forwards its tensor argument to another block.60 61#map0 = affine_map<() -> ()>62 63#attrs = {64 indexing_maps = [#map0, #map0, #map0],65 iterator_types = []66}67 68func.func @main() -> (tensor<i32>) attributes {} {69 %c0 = arith.constant 0 : i3270 %0 = tensor.from_elements %c0 : tensor<i32>71 %c10 = arith.constant 10 : i3272 %1 = tensor.from_elements %c10 : tensor<i32>73 cf.br ^bb1(%0 : tensor<i32>)74 75^bb1(%2: tensor<i32>): // 2 preds: ^bb0, ^bb276 %3 = tensor.empty() : tensor<i1>77 %4 = linalg.generic #attrs78 ins(%2, %1 : tensor<i32>, tensor<i32>)79 outs(%3 : tensor<i1>) {80 ^bb0(%arg0: i32, %arg1: i32, %arg2: i1):81 %8 = arith.cmpi slt, %arg0, %arg1 : i3282 linalg.yield %8 : i183 } -> tensor<i1>84 %5 = tensor.extract %4[] : tensor<i1>85 cf.cond_br %5, ^bb2(%2 : tensor<i32>), ^bb3(%2 : tensor<i32>)86 87^bb2(%6: tensor<i32>): // pred: ^bb188 %7 = tensor.empty() : tensor<i32>89 %8 = linalg.generic #attrs90 ins(%6, %6 : tensor<i32>, tensor<i32>)91 outs(%7 : tensor<i32>) {92 ^bb0(%arg0: i32, %arg1: i32, %arg2: i32):93 %9 = arith.addi %arg0, %arg1 : i3294 linalg.yield %9 : i3295 } -> tensor<i32>96 cf.br ^bb3(%8 : tensor<i32>)97 98^bb3(%10: tensor<i32>): // pred: ^bb199 cf.br ^bb4(%10 : tensor<i32>)100 101^bb4(%11: tensor<i32>): // pred: ^bb1102 return %11 : tensor<i32>103}104 105// CHECK-LABEL: func @main()106// CHECK-DAG: %[[cst:.*]] = arith.constant dense<0>107// CHECK-DAG: arith.constant true108// CHECK: cf.br ^[[bb1:.*]]109// CHECK-NEXT: ^[[bb1:.*]]:110// CHECK-NEXT: cf.cond_br %{{.*}}, ^[[bb2:.*]], ^bb3111// CHECK-NEXT: ^[[bb2]]:112// CHECK-NEXT: cf.br ^[[bb3:.*]]113// CHECK-NEXT: ^[[bb3]]:114// CHECK-NEXT: cf.br ^[[bb4:.*]]115// CHECK-NEXT: ^[[bb4]]:116// CHECK-NEXT: return %[[cst]]117// CHECK-NEXT: }118 119// -----120 121#map0 = affine_map<() -> ()>122 123#attrs = {124 indexing_maps = [#map0, #map0, #map0],125 iterator_types = []126}127 128func.func @main() -> (tensor<i32>) attributes {} {129 %c0 = arith.constant 0 : i32130 %0 = tensor.from_elements %c0 : tensor<i32>131 %c10 = arith.constant 10 : i32132 %1 = tensor.from_elements %c10 : tensor<i32>133 cf.br ^bb1(%0 : tensor<i32>)134 135^bb1(%2: tensor<i32>): // 2 preds: ^bb0, ^bb2136 %3 = tensor.empty() : tensor<i1>137 %4 = linalg.generic #attrs138 ins(%2, %1 : tensor<i32>, tensor<i32>)139 outs(%3 : tensor<i1>) {140 ^bb0(%arg0: i32, %arg1: i32, %arg2: i1):141 %8 = arith.cmpi slt, %arg0, %arg1 : i32142 linalg.yield %8 : i1143 } -> tensor<i1>144 %5 = tensor.extract %4[] : tensor<i1>145 // This cf.cond_br intentionally has bb2 as it's target for both branches. This146 // is to make sure that the "forward phase" of the cost-model correctly adds147 // the users of a block argument (in this case bb2's argument) to the work148 // list.149 cf.cond_br %5, ^bb2(%2 : tensor<i32>), ^bb2(%2 : tensor<i32>)150 151^bb2(%6: tensor<i32>): // pred: ^bb1152 %12 = tensor.from_elements %c10 : tensor<i32>153 %7 = tensor.empty() : tensor<i32>154 %8 = linalg.generic #attrs155 ins(%6, %12 : tensor<i32>, tensor<i32>)156 outs(%7 : tensor<i32>) {157 ^bb0(%arg0: i32, %arg1: i32, %arg2: i32):158 %9 = arith.addi %arg0, %arg1 : i32159 linalg.yield %9 : i32160 } -> tensor<i32>161 cf.br ^bb3(%8 : tensor<i32>)162 163^bb3(%10: tensor<i32>): // pred: ^bb1164 return %10 : tensor<i32>165}166 167// CHECK-LABEL: func @main()168// CHECK-DAG: %[[cst:.*]] = arith.constant dense<10>169// CHECK-DAG: arith.constant true170// CHECK: cf.br ^[[bb1:.*]]171// CHECK-NEXT: ^[[bb1]]:172// CHECK-NEXT: cf.cond_br %{{.*}}, ^[[bb2:.*]], ^bb2173// CHECK-NEXT: ^[[bb2]]174// CHECK-NEXT: cf.br ^[[bb3:.*]]175// CHECK-NEXT: ^[[bb3]]176// CHECK-NEXT: return %[[cst]]177// CHECK-NEXT: }178