73 lines · plain
1// RUN: mlir-opt %s -linalg-inline-scalar-operands -split-input-file | FileCheck %s2 3// CHECK: #[[MAP:.*]] = affine_map<(d0) -> (d0)>4#map2 = affine_map<(d0) -> (d0)>5#map3 = affine_map<(d0) -> ()>6 7// CHECK: func @inline_zerod(%[[ARG:.*]]: tensor<4xf32>, %[[SCALAR:.*]]: tensor<f32>)8func.func @inline_zerod(%arg0: tensor<4xf32>, %scalar: tensor<f32>) -> tensor<4xf32> {9 %0 = tensor.empty() : tensor<4xf32>10 // CHECK: linalg.generic {indexing_maps = [#[[MAP]], #[[MAP]]],11 // CHECK-SAME: iterator_types = ["parallel"]} ins(%[[ARG]] : tensor<4xf32>)12 %1 = linalg.generic {indexing_maps = [#map2, #map3, #map2],13 iterator_types = ["parallel"]}14 ins(%arg0, %scalar : tensor<4xf32>, tensor<f32>)15 outs(%0 : tensor<4xf32>) {16 // CHECK: ^bb0(%{{.*}}: f32, %{{.*}}: f32)17 ^bb0(%arg1: f32, %arg2: f32, %arg3: f32):18 // CHECK: tensor.extract %[[SCALAR]][]19 %2 = arith.divf %arg1, %arg2 : f3220 linalg.yield %2 : f3221 } -> tensor<4xf32>22 return %1 : tensor<4xf32>23}24 25// -----26 27// CHECK: #[[MAP:.*]] = affine_map<(d0) -> (d0)>28#map2 = affine_map<(d0) -> (d0)>29#map3 = affine_map<(d0) -> (0)>30 31// CHECK: func @inline_oned(%[[ARG:.*]]: tensor<4xf32>, %[[SCALAR:.*]]: tensor<1xf32>)32func.func @inline_oned(%arg0: tensor<4xf32>, %scalar: tensor<1xf32>) -> tensor<4xf32> {33 // CHECK: %[[ZERO:.*]] = arith.constant 0 : index34 %0 = tensor.empty() : tensor<4xf32>35 // CHECK: linalg.generic {indexing_maps = [#[[MAP]], #[[MAP]]],36 // CHECK-SAME: iterator_types = ["parallel"]} ins(%[[ARG]] : tensor<4xf32>)37 %1 = linalg.generic {indexing_maps = [#map2, #map3, #map2],38 iterator_types = ["parallel"]}39 ins(%arg0, %scalar : tensor<4xf32>, tensor<1xf32>)40 outs(%0 : tensor<4xf32>) {41 // CHECK: ^bb0(%{{.*}}: f32, %{{.*}}: f32)42 ^bb0(%arg1: f32, %arg2: f32, %arg3: f32):43 // CHECK: tensor.extract %[[SCALAR]][%[[ZERO]]]44 %2 = arith.divf %arg1, %arg2 : f3245 linalg.yield %2 : f3246 } -> tensor<4xf32>47 return %1 : tensor<4xf32>48}49 50// -----51 52// CHECK: #[[MAP:.*]] = affine_map<(d0) -> (d0)>53#map2 = affine_map<(d0) -> (d0)>54#map3 = affine_map<(d0) -> ()>55 56// CHECK: func @inline_scalar(%[[ARG:.*]]: tensor<4xf32>, %[[SCALAR:.*]]: f32)57func.func @inline_scalar(%arg0: tensor<4xf32>, %scalar: f32) -> tensor<4xf32> {58 %0 = tensor.empty() : tensor<4xf32>59 // CHECK: linalg.generic {indexing_maps = [#[[MAP]], #[[MAP]]],60 // CHECK-SAME: iterator_types = ["parallel"]} ins(%[[ARG]] : tensor<4xf32>)61 %1 = linalg.generic {indexing_maps = [#map2, #map3, #map2],62 iterator_types = ["parallel"]}63 ins(%arg0, %scalar : tensor<4xf32>, f32)64 outs(%0 : tensor<4xf32>) {65 // CHECK: ^bb0(%[[IN:.*]]: f32, %[[OUT:.*]]: f32)66 ^bb0(%arg1: f32, %arg2: f32, %arg3: f32):67 // CHECK: arith.divf %[[IN]], %[[SCALAR]] : f3268 %2 = arith.divf %arg1, %arg2 : f3269 linalg.yield %2 : f3270 } -> tensor<4xf32>71 return %1 : tensor<4xf32>72}73