177 lines · plain
1// RUN: mlir-opt -pass-pipeline="builtin.module(func.func(convert-elementwise-to-linalg))" -split-input-file %s | FileCheck %s2 3// In-depth checking of the linalg.generic op for a very trivial case.4// CHECK: #[[$MAP:.*]] = affine_map<() -> ()>5// CHECK-LABEL: func @addf_rank06// CHECK-SAME: %[[ARG0:[0-9a-zA-Z]*]]: tensor<f32>7// CHECK-SAME: %[[ARG1:[0-9a-zA-Z]*]]: tensor<f32>8func.func @addf_rank0(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {9 // CHECK: %{{.*}} = linalg.generic10 // CHECK-SAME: indexing_maps = [#[[$MAP]], #[[$MAP]], #[[$MAP]]]11 // CHECK-SAME: iterator_types = []12 // CHECK-SAME: ins(%[[ARG0]], %[[ARG1]]13 // CHECK-SAME: outs(%[[ARG0]]14 // CHECK: ^bb0(%[[LHS:.*]]: f32, %[[RHS:.*]]: f32, %{{.*}}: f32):15 // CHECK: %[[YIELD:.*]] = arith.addf %[[LHS]], %[[RHS]] : f3216 // CHECK: linalg.yield %[[YIELD]] : f3217 // CHECK: } -> tensor<f32>18 %0 = arith.addf %arg0, %arg1 : tensor<f32>19 return %0 : tensor<f32>20}21 22// -----23 24// Check indexing maps and iterator types for the rank > 0 case.25// CHECK-LABEL: func @addf_rank126// CHECK-SAME: %[[ARG0:[0-9a-zA-Z]*]]: tensor<?xf32>27// CHECK-SAME: %[[ARG1:[0-9a-zA-Z]*]]: tensor<?xf32>28func.func @addf_rank1(%arg0: tensor<?xf32>, %arg1: tensor<?xf32>) -> tensor<?xf32> {29 // CHECK: linalg.generic30 // CHECK-SAME: iterator_types = ["parallel"]31 // CHECK-SAME: ins(%[[ARG0]], %[[ARG1]]32 // CHECK-SAME: outs(%[[ARG0]]33 %0 = arith.addf %arg0, %arg1 : tensor<?xf32>34 return %0 : tensor<?xf32>35}36 37// -----38 39// Check a unary op.40// CHECK-LABEL: func @exp41// CHECK-SAME: %[[ARG0:[0-9a-zA-Z]*]]: tensor<f32>42func.func @exp(%arg0: tensor<f32>) -> tensor<f32> {43 // CHECK: linalg.generic44 // CHECK-SAME: ins(%[[ARG0]]45 // CHECK-SAME: outs(%[[ARG0]]46 // CHECK: ^bb0(%[[SCALAR:.*]]: f32, %{{.*}}: f32):47 // CHECK: %[[YIELD:.*]] = math.exp %[[SCALAR]] : f3248 // CHECK: linalg.yield %[[YIELD]] : f3249 %0 = math.exp %arg0 : tensor<f32>50 return %0 : tensor<f32>51}52 53// -----54 55// Check a case with varying operand types.56// CHECK-LABEL: func @select57// CHECK-SAME: %[[ARG0:[0-9a-zA-Z]*]]: tensor<i1>58// CHECK-SAME: %[[ARG1:[0-9a-zA-Z]*]]: tensor<i32>59// CHECK-SAME: %[[ARG2:[0-9a-zA-Z]*]]: tensor<i32>60func.func @select(%arg0: tensor<i1>, %arg1: tensor<i32>, %arg2: tensor<i32>) -> tensor<i32> {61 // CHECK: linalg.generic62 // CHECK-SAME: ins(%[[ARG0]], %[[ARG1]], %[[ARG2]]63 // CHECK-SAME: outs(%[[ARG1]]64 // CHECK: ^bb0(%[[PRED:.*]]: i1, %[[TRUE_VAL:.*]]: i32, %[[FALSE_VAL:.*]]: i32, %{{.*}}: i32):65 // CHECK: arith.select %[[PRED]], %[[TRUE_VAL]], %[[FALSE_VAL]] : i3266 %0 = arith.select %arg0, %arg1, %arg2 : tensor<i1>, tensor<i32>67 return %0 : tensor<i32>68}69 70// -----71 72// Spot-check an op that requires copying attributes properly to the created scalar op.73// Also checks proper init_tensor usage.74// CHECK-LABEL: func @cmpf(75// CHECK-SAME: %[[ARG0:[0-9a-zA-Z]*]]: tensor<f32>76// CHECK-SAME: %[[ARG1:[0-9a-zA-Z]*]]: tensor<f32>77func.func @cmpf(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {78 // CHECK: %[[INIT:.*]] = tensor.empty() : tensor<i1>79 // CHECK: linalg.generic80 // CHECK-SAME: ins(%[[ARG0]], %[[ARG1]]81 // CHECK-SAME: outs(%[[INIT]]82 // CHECK: ^bb0(%{{.*}}: f32, %{{.*}}: f32, %{{.*}}: i1):83 // CHECK: arith.cmpf olt, %{{.*}}, %{{.*}} : f3284 %0 = arith.cmpf olt, %arg0, %arg1 : tensor<f32>85 return %0 : tensor<i1>86}87 88// -----89 90// Check proper init_tensor usage in a mixed case.91// CHECK-LABEL: func @cmpf(92// CHECK-SAME: %[[ARG0:[0-9a-zA-Z]*]]: tensor<4x?x?x8x2x?xf32>93// CHECK-SAME: %[[ARG1:[0-9a-zA-Z]*]]: tensor<4x?x?x8x2x?xf32>94func.func @cmpf(%arg0: tensor<4x?x?x8x2x?xf32>, %arg1: tensor<4x?x?x8x2x?xf32>) -> tensor<4x?x?x8x2x?xi1> {95 // CHECK: %[[C1:.*]] = arith.constant 1 : index96 // CHECK: %[[D1:.*]] = tensor.dim %[[ARG0]], %[[C1]] : tensor<4x?x?x8x2x?xf32>97 // CHECK: %[[C2:.*]] = arith.constant 2 : index98 // CHECK: %[[D2:.*]] = tensor.dim %[[ARG0]], %[[C2]] : tensor<4x?x?x8x2x?xf32>99 // CHECK: %[[C5:.*]] = arith.constant 5 : index100 // CHECK: %[[D5:.*]] = tensor.dim %[[ARG0]], %[[C5]] : tensor<4x?x?x8x2x?xf32>101 // CHECK: %[[INIT:.*]] = tensor.empty(%[[D1]], %[[D2]], %[[D5]]) : tensor<4x?x?x8x2x?xi1>102 // CHECK: linalg.generic103 // CHECK-SAME: ins(%[[ARG0]], %[[ARG1]]104 // CHECK-SAME: outs(%[[INIT]]105 // CHECK: ^bb0(%{{.*}}: f32, %{{.*}}: f32, %{{.*}}: i1):106 // CHECK: arith.cmpf olt, %{{.*}}, %{{.*}} : f32107 %0 = arith.cmpf olt, %arg0, %arg1 : tensor<4x?x?x8x2x?xf32>108 return %0 : tensor<4x?x?x8x2x?xi1>109}110 111// -----112 113// Check a mix of scalar and tensor input.114// CHECK: #[[$MAP1:.*]] = affine_map<(d0, d1) -> ()> 115// CHECK: #[[$MAP2:.*]] = affine_map<(d0, d1) -> (d0, d1)> 116// CHECK-LABEL: func @scalar_plus_tensor117func.func @scalar_plus_tensor(%arg0: f32, %arg1: tensor<?x?xf32>) -> tensor<?x?xf32> {118 // CHECK: %[[GEN:.*]] = linalg.generic119 // CHECK-SAME: iterator_types = ["parallel", "parallel"]120 // CHECK-SAME: ins(%[[S:.*]], %[[T:.*]] : f32, tensor<?x?xf32>)121 // CHECK-SAME: outs(%[[T]] : tensor<?x?xf32>)122 // CHECK: ^bb0(%[[SB:.*]]: f32, %[[TB:.*]]: f32, %[[OB:.*]]: f32):123 // CHECK: "test.elementwise_mappable"(%[[SB]], %[[TB]]) : (f32, f32) -> f32124 // CHECK: linalg.yield {{.*}} : f32125 // CHECK: } -> tensor<?x?xf32>126 %0 = "test.elementwise_mappable"(%arg0, %arg1)127 : (f32, tensor<?x?xf32>) -> tensor<?x?xf32>128 return %0 : tensor<?x?xf32>129}130 131// -----132// This test exercises the case where an elementwise op has two scalar-like133// operands and one ranked tensor operand. In this example, we chain two134// `test.elementwise_mappable` calls:135// %0 = f(%s1, %t)136// %1 = f(%s2, %0)137// CHECK-DAG: #[[$SC2:[A-Za-z0-9_]+]] = affine_map<(d0, d1) -> ()>138// CHECK-DAG: #[[$ID2:[A-Za-z0-9_]+]] = affine_map<(d0, d1) -> (d0, d1)>139// CHECK-LABEL: func @scalar_tensor_scalar140func.func @scalar_tensor_scalar(%s1: f32, %t: tensor<?x?xf32>, %s2: f32) -> tensor<?x?xf32> {141 // First generic.142 // CHECK: %[[GEN0:.*]] = linalg.generic143 // CHECK-SAME: indexing_maps = [#[[$SC2]], #[[$ID2]], #[[$ID2]]]144 // CHECK-SAME: iterator_types = ["parallel", "parallel"]145 // CHECK-SAME: ins(%[[S1:[^,]+]], %[[T0:[^)]*]] : f32, tensor<?x?xf32>)146 // CHECK-SAME: outs(%[[T0]] : tensor<?x?xf32>)147 // CHECK: ^bb0(%[[S1E:.*]]: f32, %[[T0E:.*]]: f32, %[[O0E:.*]]: f32):148 // CHECK: %[[APPLY0:.*]] = "test.elementwise_mappable"(%[[S1E]], %[[T0E]]) : (f32, f32) -> f32149 // CHECK: linalg.yield %[[APPLY0]] : f32150 // CHECK: } -> tensor<?x?xf32>151 152 // Second generic.153 // CHECK: %[[GEN1:.*]] = linalg.generic154 // CHECK-SAME: indexing_maps = [#[[$SC2]], #[[$ID2]], #[[$ID2]]]155 // CHECK-SAME: iterator_types = ["parallel", "parallel"]156 // CHECK-SAME: ins(%[[S2:[^,]+]], %[[GEN0]] : f32, tensor<?x?xf32>)157 // CHECK-SAME: outs(%[[GEN0]] : tensor<?x?xf32>)158 // CHECK: ^bb0(%[[S2E:.*]]: f32, %[[G0E:.*]]: f32, %[[O1E:.*]]: f32):159 // CHECK: %[[APPLY1:.*]] = "test.elementwise_mappable"(%[[S2E]], %[[G0E]]) : (f32, f32) -> f32160 // CHECK: linalg.yield %[[APPLY1]] : f32161 // CHECK: } -> tensor<?x?xf32>162 // CHECK: return %[[GEN1]] : tensor<?x?xf32>163 %0 = "test.elementwise_mappable"(%s1, %t)164 : (f32, tensor<?x?xf32>) -> tensor<?x?xf32>165 %1 = "test.elementwise_mappable"(%s2, %0)166 : (f32, tensor<?x?xf32>) -> tensor<?x?xf32>167 return %1 : tensor<?x?xf32>168}169 170// ----171// CHECK-LABEL: func @negative_scalar_only_eltwise172// CHECK-NOT: linalg173func.func @negative_scalar_only_eltwise(%a: f32, %b: f32) -> f32 {174 %0 = arith.addf %a, %b : f32175 return %0 : f32176}177