1// RUN: mlir-opt %s -pass-pipeline="builtin.module(func.func(linalg-detensorize{aggressive-mode}))" | FileCheck %s -check-prefix=DET-ALL 2// RUN: mlir-opt %s -pass-pipeline="builtin.module(func.func(linalg-detensorize))" | FileCheck %s -check-prefix=DET-CF 3 4#map0 = affine_map<() -> ()> 5 6#attrs = { 7 indexing_maps = [#map0, #map0, #map0], 8 iterator_types = [] 9} 10 11func.func @main(%farg0: tensor<i32>, %farg1: tensor<i32>) -> tensor<i32> attributes {} { 12 cf.br ^bb1(%farg0 : tensor<i32>) 13 14^bb1(%0: tensor<i32>): // 2 preds: ^bb0, ^bb2 15 %1 = tensor.empty() : tensor<i1> 16 %2 = linalg.generic #attrs 17 ins(%0, %farg1 : tensor<i32>, tensor<i32>) 18 outs(%1 : tensor<i1>) { 19 ^bb0(%arg0: i32, %arg1: i32, %arg2: i1): 20 %8 = arith.cmpi slt, %arg0, %arg1 : i32 21 linalg.yield %8 : i1 22 } -> tensor<i1> 23 %3 = tensor.extract %2[] : tensor<i1> 24 cf.cond_br %3, ^bb2(%0 : tensor<i32>), ^bb3(%0 : tensor<i32>) 25 26^bb2(%4: tensor<i32>): // pred: ^bb1 27 %5 = tensor.empty() : tensor<i32> 28 %6 = linalg.generic #attrs 29 ins(%4, %4 : tensor<i32>, tensor<i32>) 30 outs(%5 : tensor<i32>) { 31 ^bb0(%arg0: i32, %arg1: i32, %arg2: i32): 32 %8 = arith.addi %arg0, %arg1 : i32 33 linalg.yield %8 : i32 34 } -> tensor<i32> 35 cf.br ^bb1(%6 : tensor<i32>) 36 37^bb3(%7: tensor<i32>): // pred: ^bb1 38 return %7 : tensor<i32> 39} 40 41// Test aggresively detensoring all detensorable ops. 42// 43// DET-ALL-LABEL: func @main 44// DET-ALL-SAME: (%{{.*}}: tensor<i32>, %{{.*}}: tensor<i32>) 45// DET-ALL: tensor.extract {{.*}} 46// DET-ALL: cf.br ^[[bb1:.*]](%{{.*}} : i32) 47// DET-ALL: ^[[bb1]](%{{.*}}: i32) 48// DET-ALL: arith.cmpi slt, {{.*}} 49// DET-ALL: cf.cond_br {{.*}}, ^[[bb2:.*]], ^[[bb3:.*]] 50// DET-ALL: ^[[bb2]] 51// DET-ALL: arith.addi {{.*}} 52// DET-ALL: cf.br ^[[bb1]](%{{.*}} : i32) 53// DET-ALL: ^[[bb3]]: 54// DET-ALL: tensor.from_elements {{.*}} 55// DET-ALL: return %{{.*}} : tensor<i32> 56 57// Test detensoring only ops involed in control-flow. 58// 59// DET-CF-LABEL: func @main 60// DET-CF-SAME: (%{{.*}}: tensor<i32>, %{{.*}}: tensor<i32>) 61// DET-CF: tensor.extract {{.*}} 62// DET-CF: cf.br ^[[bb1:.*]](%{{.*}} : i32) 63// DET-CF: ^[[bb1]](%{{.*}}: i32) 64// DET-CF: arith.cmpi slt, {{.*}} 65// DET-CF: cf.cond_br {{.*}}, ^[[bb2:.*]], ^[[bb3:.*]] 66// DET-CF: ^[[bb2]]: 67// DET-CF: arith.addi {{.*}} 68// DET-CF: cf.br ^[[bb1]](%{{.*}} : i32) 69// DET-CF: ^[[bb3]]: 70// DET-CF: tensor.from_elements %{{.*}} : tensor<i32> 71// DET-CF: return %{{.*}} : tensor<i32> 72