59 lines · plain
1// RUN: mlir-opt %s -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() -> () attributes {} {11 %c0 = arith.constant 0 : i3212 %0 = tensor.from_elements %c0 : tensor<1xi32>13 %reshaped0 = tensor.collapse_shape %0 [] : tensor<1xi32> into tensor<i32>14 %c10 = arith.constant 10 : i3215 %1 = tensor.from_elements %c10 : tensor<1xi32>16 %reshaped1 = tensor.collapse_shape %1 [] : tensor<1xi32> into tensor<i32>17 cf.br ^bb1(%reshaped0 : tensor<i32>)18 19^bb1(%2: tensor<i32>): // 2 preds: ^bb0, ^bb220 %3 = tensor.empty() : tensor<i1>21 %4 = linalg.generic #attrs22 ins(%2, %reshaped1 : tensor<i32>, tensor<i32>)23 outs(%3 : tensor<i1>) {24 ^bb0(%arg0: i32, %arg1: i32, %arg2: i1):25 %8 = arith.cmpi slt, %arg0, %arg1 : i3226 linalg.yield %8 : i127 } -> tensor<i1>28 %5 = tensor.extract %4[] : tensor<i1>29 cf.cond_br %5, ^bb2(%2 : tensor<i32>), ^bb330 31^bb2(%6: tensor<i32>): // pred: ^bb132 %7 = tensor.empty() : tensor<i32>33 %8 = linalg.generic #attrs34 ins(%6, %6 : tensor<i32>, tensor<i32>)35 outs(%7 : tensor<i32>) {36 ^bb0(%arg0: i32, %arg1: i32, %arg2: i32):37 %9 = arith.addi %arg0, %arg1 : i3238 linalg.yield %9 : i3239 } -> tensor<i32>40 cf.br ^bb1(%8 : tensor<i32>)41 42^bb3: // pred: ^bb143 return44}45 46// CHECK-LABEL: func @main47// CHECK-DAG: arith.constant 0 : i3248// CHECK-DAG: arith.constant 1049// CHECK-NEXT: cf.br ^[[bb1:.*]](%{{.*}} : i32)50// CHECK-NEXT: ^[[bb1]](%{{.*}}: i32)51// CHECK-NEXT: %{{.*}} = arith.cmpi slt, %{{.*}}, %{{.*}}52// CHECK-NEXT: cf.cond_br %{{.*}}, ^[[bb2:.*]], ^[[bb3:.*]]53// CHECK-NEXT: ^[[bb2]]54// CHECK-NEXT: %{{.*}} = arith.addi %{{.*}}, %{{.*}}55// CHECK-NEXT: cf.br ^[[bb1]](%{{.*}} : i32)56// CHECK-NEXT: ^[[bb3]]:57// CHECK-NEXT: return58// CHECK-NEXT: }59