| // RUN: mlir-opt %s -split-input-file -allow-unregistered-dialect -linalg-detensorize | FileCheck %s |
| |
| #map0 = affine_map<() -> ()> |
| |
| #attrs = { |
| indexing_maps = [#map0, #map0, #map0], |
| iterator_types = [] |
| } |
| |
| func @main() -> (tensor<i32>) attributes {} { |
| %c0 = arith.constant 0 : i32 |
| %0 = tensor.from_elements %c0 : tensor<1xi32> |
| %reshaped0 = linalg.tensor_collapse_shape %0 [] : tensor<1xi32> into tensor<i32> |
| %c10 = arith.constant 10 : i32 |
| %1 = tensor.from_elements %c10 : tensor<1xi32> |
| %reshaped1 = linalg.tensor_collapse_shape %1 [] : tensor<1xi32> into tensor<i32> |
| br ^bb1(%reshaped0 : tensor<i32>) |
| |
| ^bb1(%2: tensor<i32>): // 2 preds: ^bb0, ^bb2 |
| %3 = linalg.init_tensor [] : tensor<i1> |
| %4 = linalg.generic #attrs |
| ins(%2, %reshaped1 : tensor<i32>, tensor<i32>) |
| outs(%3 : tensor<i1>) { |
| ^bb0(%arg0: i32, %arg1: i32, %arg2: i1): // no predecessors |
| %8 = arith.cmpi slt, %arg0, %arg1 : i32 |
| linalg.yield %8 : i1 |
| } -> tensor<i1> |
| %5 = tensor.extract %4[] : tensor<i1> |
| cond_br %5, ^bb2(%2 : tensor<i32>), ^bb3(%2 : tensor<i32>) |
| |
| ^bb2(%6: tensor<i32>): // pred: ^bb1 |
| %7 = linalg.init_tensor [] : tensor<i32> |
| %8 = linalg.generic #attrs |
| ins(%6, %6 : tensor<i32>, tensor<i32>) |
| outs(%7 : tensor<i32>) { |
| ^bb0(%arg0: i32, %arg1: i32, %arg2: i32): // no predecessors |
| %9 = arith.addi %arg0, %arg1 : i32 |
| linalg.yield %9 : i32 |
| } -> tensor<i32> |
| br ^bb3(%8 : tensor<i32>) |
| |
| ^bb3(%10: tensor<i32>): // pred: ^bb1 |
| return %10 : tensor<i32> |
| } |
| |
| // CHECK-LABEL: func @main() |
| // CHECK-NEXT: arith.constant 0 |
| // CHECK-NEXT: arith.constant 10 |
| // CHECK-NEXT: br ^[[bb1:.*]](%{{.*}}: i32) |
| // CHECK-NEXT: ^[[bb1]](%{{.*}}: i32): |
| // CHECK-NEXT: arith.cmpi slt, %{{.*}}, %{{.*}} |
| // CHECK-NEXT: cond_br %{{.*}}, ^[[bb2:.*]](%{{.*}} : i32), ^bb3(%{{.*}} : i32) |
| // CHECK-NEXT: ^[[bb2]](%{{.*}}: i32) |
| // CHECK-NEXT: arith.addi %{{.*}}, %{{.*}} |
| // CHECK-NEXT: br ^[[bb3:.*]](%{{.*}} : i32) |
| // CHECK-NEXT: ^[[bb3]](%{{.*}}: i32) |
| // CHECK-NEXT: tensor.from_elements %{{.*}} : tensor<1xi32> |
| // CHECK-NEXT: linalg.tensor_collapse_shape %{{.*}} [] : tensor<1xi32> into tensor<i32> |
| // CHECK-NEXT: return %{{.*}} |
| // CHECK-NEXT: } |
| |
| // ----- |
| |
| // Similar to the above test with one change: one of the block after the |
| // if-condition passes/forwards its tensor argument to another block. |
| |
| #map0 = affine_map<() -> ()> |
| |
| #attrs = { |
| indexing_maps = [#map0, #map0, #map0], |
| iterator_types = [] |
| } |
| |
| func @main() -> (tensor<i32>) attributes {} { |
| %c0 = arith.constant 0 : i32 |
| %0 = tensor.from_elements %c0 : tensor<1xi32> |
| %reshaped0 = linalg.tensor_collapse_shape %0 [] : tensor<1xi32> into tensor<i32> |
| %c10 = arith.constant 10 : i32 |
| %1 = tensor.from_elements %c10 : tensor<1xi32> |
| %reshaped1 = linalg.tensor_collapse_shape %1 [] : tensor<1xi32> into tensor<i32> |
| br ^bb1(%reshaped0 : tensor<i32>) |
| |
| ^bb1(%2: tensor<i32>): // 2 preds: ^bb0, ^bb2 |
| %3 = linalg.init_tensor [] : tensor<i1> |
| %4 = linalg.generic #attrs |
| ins(%2, %reshaped1 : tensor<i32>, tensor<i32>) |
| outs(%3 : tensor<i1>) { |
| ^bb0(%arg0: i32, %arg1: i32, %arg2: i1): // no predecessors |
| %8 = arith.cmpi slt, %arg0, %arg1 : i32 |
| linalg.yield %8 : i1 |
| } -> tensor<i1> |
| %5 = tensor.extract %4[] : tensor<i1> |
| cond_br %5, ^bb2(%2 : tensor<i32>), ^bb3(%2 : tensor<i32>) |
| |
| ^bb2(%6: tensor<i32>): // pred: ^bb1 |
| %7 = linalg.init_tensor [] : tensor<i32> |
| %8 = linalg.generic #attrs |
| ins(%6, %6 : tensor<i32>, tensor<i32>) |
| outs(%7 : tensor<i32>) { |
| ^bb0(%arg0: i32, %arg1: i32, %arg2: i32): // no predecessors |
| %9 = arith.addi %arg0, %arg1 : i32 |
| linalg.yield %9 : i32 |
| } -> tensor<i32> |
| br ^bb3(%8 : tensor<i32>) |
| |
| ^bb3(%10: tensor<i32>): // pred: ^bb1 |
| br ^bb4(%10 : tensor<i32>) |
| |
| ^bb4(%11: tensor<i32>): // pred: ^bb1 |
| return %11 : tensor<i32> |
| } |
| |
| // CHECK-LABEL: func @main() |
| // CHECK-NEXT: arith.constant 0 |
| // CHECK-NEXT: arith.constant 10 |
| // CHECK-NEXT: br ^[[bb1:.*]](%{{.*}}: i32) |
| // CHECK-NEXT: ^[[bb1]](%{{.*}}: i32): |
| // CHECK-NEXT: arith.cmpi slt, %{{.*}}, %{{.*}} |
| // CHECK-NEXT: cond_br %{{.*}}, ^[[bb2:.*]](%{{.*}} : i32), ^bb3(%{{.*}} : i32) |
| // CHECK-NEXT: ^[[bb2]](%{{.*}}: i32) |
| // CHECK-NEXT: arith.addi %{{.*}}, %{{.*}} |
| // CHECK-NEXT: br ^[[bb3:.*]](%{{.*}} : i32) |
| // CHECK-NEXT: ^[[bb3]](%{{.*}}: i32) |
| // CHECK-NEXT: br ^[[bb4:.*]](%{{.*}} : i32) |
| // CHECK-NEXT: ^[[bb4]](%{{.*}}: i32) |
| // CHECK-NEXT: tensor.from_elements %{{.*}} : tensor<1xi32> |
| // CHECK-NEXT: linalg.tensor_collapse_shape %{{.*}} [] : tensor<1xi32> into tensor<i32> |
| // CHECK-NEXT: return %{{.*}} |
| // CHECK-NEXT: } |
| |
| // ----- |
| |
| #map0 = affine_map<() -> ()> |
| |
| #attrs = { |
| indexing_maps = [#map0, #map0, #map0], |
| iterator_types = [] |
| } |
| |
| func @main() -> (tensor<i32>) attributes {} { |
| %c0 = arith.constant 0 : i32 |
| %0 = tensor.from_elements %c0 : tensor<1xi32> |
| %reshaped0 = linalg.tensor_collapse_shape %0 [] : tensor<1xi32> into tensor<i32> |
| %c10 = arith.constant 10 : i32 |
| %1 = tensor.from_elements %c10 : tensor<1xi32> |
| %reshaped1 = linalg.tensor_collapse_shape %1 [] : tensor<1xi32> into tensor<i32> |
| br ^bb1(%reshaped0 : tensor<i32>) |
| |
| ^bb1(%2: tensor<i32>): // 2 preds: ^bb0, ^bb2 |
| %3 = linalg.init_tensor [] : tensor<i1> |
| %4 = linalg.generic #attrs |
| ins(%2, %reshaped1 : tensor<i32>, tensor<i32>) |
| outs(%3 : tensor<i1>) { |
| ^bb0(%arg0: i32, %arg1: i32, %arg2: i1): // no predecessors |
| %8 = arith.cmpi slt, %arg0, %arg1 : i32 |
| linalg.yield %8 : i1 |
| } -> tensor<i1> |
| %5 = tensor.extract %4[] : tensor<i1> |
| // This cond_br intentionally has bb2 as it's target for both branches. This |
| // is to make sure that the "forward phase" of the cost-model correctly adds |
| // the users of a block argument (in this case bb2's argument) to the work |
| // list. |
| cond_br %5, ^bb2(%2 : tensor<i32>), ^bb2(%2 : tensor<i32>) |
| |
| ^bb2(%6: tensor<i32>): // pred: ^bb1 |
| %12 = tensor.from_elements %c10 : tensor<1xi32> |
| %reshaped12 = linalg.tensor_collapse_shape %12 [] : tensor<1xi32> into tensor<i32> |
| %7 = linalg.init_tensor [] : tensor<i32> |
| %8 = linalg.generic #attrs |
| ins(%6, %reshaped12 : tensor<i32>, tensor<i32>) |
| outs(%7 : tensor<i32>) { |
| ^bb0(%arg0: i32, %arg1: i32, %arg2: i32): // no predecessors |
| %9 = arith.addi %arg0, %arg1 : i32 |
| linalg.yield %9 : i32 |
| } -> tensor<i32> |
| br ^bb3(%8 : tensor<i32>) |
| |
| ^bb3(%10: tensor<i32>): // pred: ^bb1 |
| return %10 : tensor<i32> |
| } |
| |
| // CHECK-LABEL: func @main() |
| // CHECK-NEXT: arith.constant 0 |
| // CHECK-NEXT: arith.constant 10 |
| // CHECK-NEXT: br ^[[bb1:.*]](%{{.*}}: i32) |
| // CHECK-NEXT: ^[[bb1]](%{{.*}}: i32): |
| // CHECK-NEXT: arith.cmpi slt, %{{.*}}, %{{.*}} |
| // CHECK-NEXT: cond_br %{{.*}}, ^[[bb2:.*]](%{{.*}} : i32), ^bb2(%{{.*}} : i32) |
| // CHECK-NEXT: ^[[bb2]](%{{.*}}: i32) |
| // CHECK-NEXT: arith.addi %{{.*}}, %{{.*}} |
| // CHECK-NEXT: br ^[[bb3:.*]](%{{.*}} : i32) |
| // CHECK-NEXT: ^[[bb3]](%{{.*}}: i32) |
| // CHECK-NEXT: tensor.from_elements %{{.*}} : tensor<1xi32> |
| // CHECK-NEXT: linalg.tensor_collapse_shape %{{.*}} [] : tensor<1xi32> into tensor<i32> |
| // CHECK-NEXT: return %{{.*}} |
| // CHECK-NEXT: } |