1// RUN: mlir-opt %s -split-input-file -allow-unregistered-dialect -pass-pipeline="builtin.module(func.func(linalg-detensorize))" | FileCheck %s 2 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 : i32 12 %0 = tensor.from_elements %c0 : tensor<i32> 13 %c10 = arith.constant 10 : i32 14 %1 = tensor.from_elements %c10 : tensor<i32> 15 cf.br ^bb1(%0 : tensor<i32>) 16 17^bb1(%2: tensor<i32>): // 2 preds: ^bb0, ^bb2 18 %3 = tensor.empty() : tensor<i1> 19 %4 = linalg.generic #attrs 20 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 : i32 24 linalg.yield %8 : i1 25 } -> 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: ^bb1 30 %7 = tensor.empty() : tensor<i32> 31 %8 = linalg.generic #attrs 32 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 : i32 36 linalg.yield %9 : i32 37 } -> tensor<i32> 38 cf.br ^bb3(%8 : tensor<i32>) 39 40^bb3(%10: tensor<i32>): // pred: ^bb1 41 return %10 : tensor<i32> 42} 43 44// CHECK-LABEL: func @main() 45// CHECK-DAG: %[[cst:.*]] = arith.constant dense<0> 46// CHECK-DAG: arith.constant true 47// CHECK: cf.br 48// CHECK-NEXT: ^[[bb1:.*]]: 49// CHECK-NEXT: cf.cond_br %{{.*}}, ^[[bb2:.*]], ^bb3 50// 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 the 59// 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 : i32 70 %0 = tensor.from_elements %c0 : tensor<i32> 71 %c10 = arith.constant 10 : i32 72 %1 = tensor.from_elements %c10 : tensor<i32> 73 cf.br ^bb1(%0 : tensor<i32>) 74 75^bb1(%2: tensor<i32>): // 2 preds: ^bb0, ^bb2 76 %3 = tensor.empty() : tensor<i1> 77 %4 = linalg.generic #attrs 78 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 : i32 82 linalg.yield %8 : i1 83 } -> 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: ^bb1 88 %7 = tensor.empty() : tensor<i32> 89 %8 = linalg.generic #attrs 90 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 : i32 94 linalg.yield %9 : i32 95 } -> tensor<i32> 96 cf.br ^bb3(%8 : tensor<i32>) 97 98^bb3(%10: tensor<i32>): // pred: ^bb1 99 cf.br ^bb4(%10 : tensor<i32>) 100 101^bb4(%11: tensor<i32>): // pred: ^bb1 102 return %11 : tensor<i32> 103} 104 105// CHECK-LABEL: func @main() 106// CHECK-DAG: %[[cst:.*]] = arith.constant dense<0> 107// CHECK-DAG: arith.constant true 108// CHECK: cf.br ^[[bb1:.*]] 109// CHECK-NEXT: ^[[bb1:.*]]: 110// CHECK-NEXT: cf.cond_br %{{.*}}, ^[[bb2:.*]], ^bb3 111// 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 : i32 130 %0 = tensor.from_elements %c0 : tensor<i32> 131 %c10 = arith.constant 10 : i32 132 %1 = tensor.from_elements %c10 : tensor<i32> 133 cf.br ^bb1(%0 : tensor<i32>) 134 135^bb1(%2: tensor<i32>): // 2 preds: ^bb0, ^bb2 136 %3 = tensor.empty() : tensor<i1> 137 %4 = linalg.generic #attrs 138 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 : i32 142 linalg.yield %8 : i1 143 } -> tensor<i1> 144 %5 = tensor.extract %4[] : tensor<i1> 145 // This cf.cond_br intentionally has bb2 as it's target for both branches. This 146 // is to make sure that the "forward phase" of the cost-model correctly adds 147 // the users of a block argument (in this case bb2's argument) to the work 148 // list. 149 cf.cond_br %5, ^bb2(%2 : tensor<i32>), ^bb2(%2 : tensor<i32>) 150 151^bb2(%6: tensor<i32>): // pred: ^bb1 152 %12 = tensor.from_elements %c10 : tensor<i32> 153 %7 = tensor.empty() : tensor<i32> 154 %8 = linalg.generic #attrs 155 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 : i32 159 linalg.yield %9 : i32 160 } -> tensor<i32> 161 cf.br ^bb3(%8 : tensor<i32>) 162 163^bb3(%10: tensor<i32>): // pred: ^bb1 164 return %10 : tensor<i32> 165} 166 167// CHECK-LABEL: func @main() 168// CHECK-DAG: %[[cst:.*]] = arith.constant dense<10> 169// CHECK-DAG: arith.constant true 170// CHECK: cf.br ^[[bb1:.*]] 171// CHECK-NEXT: ^[[bb1]]: 172// CHECK-NEXT: cf.cond_br %{{.*}}, ^[[bb2:.*]], ^bb2 173// CHECK-NEXT: ^[[bb2]] 174// CHECK-NEXT: cf.br ^[[bb3:.*]] 175// CHECK-NEXT: ^[[bb3]] 176// CHECK-NEXT: return %[[cst]] 177// CHECK-NEXT: } 178