1// RUN: mlir-opt %s -allow-unregistered-dialect -pass-pipeline="builtin.module(func.func(linalg-detensorize{aggressive-mode}))" | FileCheck %s 2 3#map = affine_map<() -> ()> 4 5func.func @detensor_simple(%arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> attributes {iree.module.export} { 6 %0 = tensor.empty() : tensor<f32> 7 %1 = linalg.generic {indexing_maps = [#map, #map, #map], iterator_types = []} 8 ins(%arg1, %arg2 : tensor<f32>, tensor<f32>) 9 outs(%0 : tensor<f32>) { 10 ^bb0(%arg3: f32, %arg4: f32, %arg5: f32): 11 %2 = arith.addf %arg3, %arg4 : f32 12 linalg.yield %2 : f32 13 } -> tensor<f32> 14 return %1: tensor<f32> 15} 16// CHECK-LABEL: func @detensor_simple 17// CHECK-SAME: (%[[arg1:.*]]: tensor<f32>, %[[arg2:.*]]: tensor<f32>) 18// CHECK-DAG: %[[arg1_val:.*]] = tensor.extract %[[arg1]] 19// CHECK-DAG: %[[arg2_val:.*]] = tensor.extract %[[arg2]] 20// CHECK: %[[detensored_res:.*]] = arith.addf %[[arg1_val]], %[[arg2_val]] 21// CHECK: %[[new_tensor_res:.*]] = tensor.from_elements %[[detensored_res]] 22// CHECK: return %[[new_tensor_res]] 23 24func.func @detensor_op_sequence(%arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> attributes {iree.module.export} { 25 %0 = tensor.empty() : tensor<f32> 26 %1 = linalg.generic {indexing_maps = [#map, #map, #map], iterator_types = []} 27 ins(%arg1, %arg2 : tensor<f32>, tensor<f32>) 28 outs(%0 : tensor<f32>) { 29 ^bb0(%arg3: f32, %arg4: f32, %arg5: f32): 30 %2 = arith.addf %arg3, %arg4 : f32 31 linalg.yield %2 : f32 32 } -> tensor<f32> 33 34 %3 = tensor.empty() : tensor<f32> 35 %4 = linalg.generic {indexing_maps = [#map, #map, #map], iterator_types = []} 36 ins(%arg1, %1 : tensor<f32>, tensor<f32>) 37 outs(%3 : tensor<f32>) { 38 ^bb0(%arg3: f32, %arg4: f32, %arg5: f32): 39 %5 = arith.mulf %arg3, %arg4 : f32 40 linalg.yield %5 : f32 41 } -> tensor<f32> 42 43 %6 = tensor.empty() : tensor<f32> 44 %7 = linalg.generic {indexing_maps = [#map, #map, #map], iterator_types = []} 45 ins(%1, %4 : tensor<f32>, tensor<f32>) 46 outs(%6 : tensor<f32>) { 47 ^bb0(%arg3: f32, %arg4: f32, %arg5: f32): 48 %5 = arith.divf %arg3, %arg4 : f32 49 linalg.yield %5 : f32 50 } -> tensor<f32> 51 52 return %7: tensor<f32> 53} 54// CHECK-LABEL: func @detensor_op_sequence 55// CHECK-SAME: (%[[arg1:.*]]: tensor<f32>, %[[arg2:.*]]: tensor<f32>) 56// CHECK-DAG: %[[arg1_val:.*]] = tensor.extract %[[arg1]] 57// CHECK-DAG: %[[arg2_val:.*]] = tensor.extract %[[arg2]] 58// CHECK: %[[detensored_res:.*]] = arith.addf %[[arg1_val]], %[[arg2_val]] 59// CHECK: %[[detensored_res2:.*]] = arith.mulf %[[arg1_val]], %[[detensored_res]] 60// CHECK: %[[detensored_res3:.*]] = arith.divf %[[detensored_res]], %[[detensored_res2]] 61// CHECK: %[[new_tensor_res:.*]] = tensor.from_elements %[[detensored_res3]] 62// CHECK: return %[[new_tensor_res]] 63 64func.func @detensor_multiple_ops(%arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> attributes {iree.module.export} { 65 %0 = tensor.empty() : tensor<f32> 66 %1 = linalg.generic {indexing_maps = [#map, #map, #map], iterator_types = []} 67 ins(%arg1, %arg2 : tensor<f32>, tensor<f32>) 68 outs(%0 : tensor<f32>) { 69 ^bb0(%arg3: f32, %arg4: f32, %arg5: f32): 70 %2 = arith.addf %arg3, %arg4 : f32 71 %3 = arith.mulf %2, %arg4 : f32 72 linalg.yield %3 : f32 73 } -> tensor<f32> 74 return %1: tensor<f32> 75} 76// CHECK-LABEL: func @detensor_multiple_ops 77// CHECK-SAME: (%[[arg1:.*]]: tensor<f32>, %[[arg2:.*]]: tensor<f32>) 78// CHECK-DAG: %[[arg1_val:.*]] = tensor.extract %[[arg1]] 79// CHECK-DAG: %[[arg2_val:.*]] = tensor.extract %[[arg2]] 80// CHECK: %[[detensored_res:.*]] = arith.addf %[[arg1_val]], %[[arg2_val]] 81// CHECK: %[[detensored_res2:.*]] = arith.mulf %[[detensored_res]], %[[arg2_val]] 82// CHECK: %[[new_tensor_res:.*]] = tensor.from_elements %[[detensored_res2]] 83// CHECK: return %[[new_tensor_res]] 84 85func.func @detensor_foreign_op(%arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> attributes {iree.module.export} { 86 %0 = tensor.empty() : tensor<f32> 87 %1 = linalg.generic {indexing_maps = [#map, #map, #map], iterator_types = []} 88 ins(%arg1, %arg2 : tensor<f32>, tensor<f32>) 89 outs(%0 : tensor<f32>) { 90 ^bb0(%arg3: f32, %arg4: f32, %arg5: f32): 91 %2 = "foreign.do_something"(%arg3, %arg4) {} : (f32, f32) -> f32 92 linalg.yield %2 : f32 93 } -> tensor<f32> 94 return %1: tensor<f32> 95} 96// CHECK-LABEL: func @detensor_foreign_op 97// CHECK-SAME: (%[[arg1:.*]]: tensor<f32>, %[[arg2:.*]]: tensor<f32>) 98// CHECK-DAG: %[[arg1_val:.*]] = tensor.extract %[[arg1]] 99// CHECK-DAG: %[[arg2_val:.*]] = tensor.extract %[[arg2]] 100// CHECK: %[[detensored_res:.*]] = "foreign.do_something"(%[[arg1_val]], %[[arg2_val]]) 101// CHECK: %[[new_tensor_res:.*]] = tensor.from_elements %[[detensored_res]] 102// CHECK: return %[[new_tensor_res]] 103