1// RUN: mlir-opt %s -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() -> () attributes {} { 11 %c0 = arith.constant 0 : i32 12 %0 = tensor.from_elements %c0 : tensor<1xi32> 13 %reshaped0 = tensor.collapse_shape %0 [] : tensor<1xi32> into tensor<i32> 14 %c10 = arith.constant 10 : i32 15 %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, ^bb2 20 %3 = tensor.empty() : tensor<i1> 21 %4 = linalg.generic #attrs 22 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 : i32 26 linalg.yield %8 : i1 27 } -> tensor<i1> 28 %5 = tensor.extract %4[] : tensor<i1> 29 cf.cond_br %5, ^bb2(%2 : tensor<i32>), ^bb3 30 31^bb2(%6: tensor<i32>): // pred: ^bb1 32 %7 = tensor.empty() : tensor<i32> 33 %8 = linalg.generic #attrs 34 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 : i32 38 linalg.yield %9 : i32 39 } -> tensor<i32> 40 cf.br ^bb1(%8 : tensor<i32>) 41 42^bb3: // pred: ^bb1 43 return 44} 45 46// CHECK-LABEL: func @main 47// CHECK-DAG: arith.constant 0 : i32 48// CHECK-DAG: arith.constant 10 49// 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: return 58// CHECK-NEXT: } 59