103 lines · plain
1// RUN: mlir-opt %s -allow-unregistered-dialect -pass-pipeline="builtin.module(func.func(linalg-detensorize{aggressive-mode}))" | FileCheck %s2 3#map = affine_map<() -> ()>4 5func.func @detensor_simple(%arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> {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 : f3212 linalg.yield %2 : f3213 } -> tensor<f32>14 return %1: tensor<f32>15}16// CHECK-LABEL: func @detensor_simple17// 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> {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 : f3231 linalg.yield %2 : f3232 } -> 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 : f3240 linalg.yield %5 : f3241 } -> 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 : f3249 linalg.yield %5 : f3250 } -> tensor<f32>51 52 return %7: tensor<f32>53}54// CHECK-LABEL: func @detensor_op_sequence55// 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> {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 : f3271 %3 = arith.mulf %2, %arg4 : f3272 linalg.yield %3 : f3273 } -> tensor<f32>74 return %1: tensor<f32>75}76// CHECK-LABEL: func @detensor_multiple_ops77// 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> {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) -> f3292 linalg.yield %2 : f3293 } -> tensor<f32>94 return %1: tensor<f32>95}96// CHECK-LABEL: func @detensor_foreign_op97// 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