166 lines · plain
1// RUN: mlir-opt %s -linalg-generalize-named-ops -split-input-file | FileCheck %s2// CHECK: #[[IDENTITY:.+]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>3//4// CHECK: @unary_exp(%[[A:.+]]: tensor<8x16x32xf32>, %[[B:.+]]: tensor<8x16x32xf32>)5// CHECK: linalg.generic6// CHECK-SAME: indexing_maps = [#[[IDENTITY]], #[[IDENTITY]]]7// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel"]8// CHECK-SAME: ins(%[[A]]9// CHECK-SAME: outs(%[[B]]10//11// CHECK: ^{{.*}}(%[[A_ARG:.+]]: f32, %[[B_ARG:.+]]: f32)12// CHECK: %[[EXP:.+]] = math.exp %[[A_ARG]] : f3213// CHECK: linalg.yield %[[EXP]] : f3214//15func.func @unary_exp(%A : tensor<8x16x32xf32>, %B: tensor<8x16x32xf32>) -> tensor<8x16x32xf32> {16 %r = linalg.elementwise17 kind=#linalg.elementwise_kind<exp>18 ins(%A : tensor<8x16x32xf32>)19 outs(%B: tensor<8x16x32xf32>) -> tensor<8x16x32xf32>20 return %r : tensor<8x16x32xf32>21}22// -----23// CHECK-DAG: #[[IDENTITY:.+]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>24// CHECK-DAG: #[[PROJECTION:.+]] = affine_map<(d0, d1, d2) -> (d2, d1)>25//26// CHECK: @unary_transpose_broadcast_tanh(%[[A:.+]]: tensor<32x16xf32>, %[[B:.+]]: tensor<8x16x32xf32>)27// CHECK: linalg.generic28// CHECK-SAME: indexing_maps = [#[[PROJECTION]], #[[IDENTITY]]]29// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel"]30// CHECK-SAME: ins(%[[A]]31// CHECK-SAME: outs(%[[B]]32//33// CHECK: ^{{.*}}(%[[A_ARG:.+]]: f32, %[[B_ARG:.+]]: f32)34// CHECK: %[[TANH:.+]] = math.tanh %[[A_ARG]] : f3235// CHECK: linalg.yield %[[TANH]] : f3236//37func.func @unary_transpose_broadcast_tanh(%A : tensor<32x16xf32>, %B: tensor<8x16x32xf32>) -> tensor<8x16x32xf32> {38 %r = linalg.elementwise39 kind=#linalg.elementwise_kind<tanh>40 indexing_maps = [affine_map<(d0, d1, d2) -> (d2, d1)>,41 affine_map<(d0, d1, d2) -> (d0, d1, d2)>]42 ins(%A : tensor<32x16xf32>)43 outs(%B: tensor<8x16x32xf32>) -> tensor<8x16x32xf32>44 return %r : tensor<8x16x32xf32>45}46// -----47// CHECK: #[[MAP:.+]] = affine_map<(d0, d1) -> (d0, d1)>48//49// CHECK: @binary_div_on_memrefs(%[[A:.+]]: memref<16x8xf32>, %[[B:.+]]: memref<16x8xf32>, %[[C:.+]]: memref<16x8xf32>)50// CHECK: linalg.generic51// CHECK-SAME: indexing_maps = [#[[MAP]], #[[MAP]], #[[MAP]]]52// CHECK-SAME: iterator_types = ["parallel", "parallel"]53// CHECK-SAME: ins(%[[A]], %[[B]]54// CHECK-SAME: outs(%[[C]]55//56// CHECK: ^{{.*}}(%[[A_ARG:.+]]: f32, %[[B_ARG:.+]]: f32, %[[C_ARG:.+]]: f32)57// CHECK: %[[DIV:.+]] = arith.divf %[[A_ARG]], %[[B_ARG]] : f3258// CHECK: linalg.yield %[[DIV]] : f3259//60func.func @binary_div_on_memrefs(%A : memref<16x8xf32>, %B: memref<16x8xf32>, %C: memref<16x8xf32>) {61 linalg.elementwise62 kind=#linalg.elementwise_kind<div>63 ins(%A, %B: memref<16x8xf32>, memref<16x8xf32>)64 outs(%C: memref<16x8xf32>)65 return66}67// -----68// CHECK: #[[MAP:.+]] = affine_map<(d0, d1) -> (d0, d1)>69//70// CHECK: @binary_mul_on_tensors(%[[A:.+]]: tensor<16x8xf32>, %[[B:.+]]: tensor<16x8xf32>, %[[C:.+]]: tensor<16x8xf32>)71// CHECK: linalg.generic72// CHECK-SAME: indexing_maps = [#[[MAP]], #[[MAP]], #[[MAP]]]73// CHECK-SAME: iterator_types = ["parallel", "parallel"]74// CHECK-SAME: ins(%[[A]], %[[B]]75// CHECK-SAME: outs(%[[C]]76//77// CHECK: ^{{.*}}(%[[A_ARG:.+]]: f32, %[[B_ARG:.+]]: f32, %[[C_ARG:.+]]: f32)78// CHECK: %[[MUL:.+]] = arith.mulf %[[A_ARG]], %[[B_ARG]] : f3279// CHECK: linalg.yield %[[MUL]] : f3280//81func.func @binary_mul_on_tensors(%A : tensor<16x8xf32>, %B: tensor<16x8xf32>, %C: tensor<16x8xf32>) -> tensor<16x8xf32> {82 %r = linalg.elementwise83 kind=#linalg.elementwise_kind<mul>84 ins(%A, %B: tensor<16x8xf32>, tensor<16x8xf32>)85 outs(%C: tensor<16x8xf32>) -> tensor<16x8xf32>86 return %r : tensor<16x8xf32>87}88// -----89// CHECK-DAG: #[[IDENTITY:.+]] = affine_map<(d0, d1) -> (d0, d1)>90// CHECK-DAG: #[[TRANSPOSE:.+]] = affine_map<(d0, d1) -> (d1, d0)>91//92// CHECK: @binary_transpose_a(%[[A:.+]]: tensor<8x16xf32>, %[[B:.+]]: tensor<16x8xf32>, %[[C:.+]]: tensor<16x8xf32>)93// CHECK: linalg.generic94// CHECK-SAME: indexing_maps = [#[[TRANSPOSE]], #[[IDENTITY]], #[[IDENTITY]]]95// CHECK-SAME: iterator_types = ["parallel", "parallel"]96// CHECK-SAME: ins(%[[A]], %[[B]]97// CHECK-SAME: outs(%[[C]]98//99// CHECK: ^{{.*}}(%[[A_ARG:.+]]: f32, %[[B_ARG:.+]]: f32, %[[C_ARG:.+]]: f32)100// CHECK: %[[SUB:.+]] = arith.subf %[[A_ARG]], %[[B_ARG]] : f32101// CHECK: linalg.yield %[[SUB]] : f32102//103func.func @binary_transpose_a(%A : tensor<8x16xf32>, %B: tensor<16x8xf32>, %C: tensor<16x8xf32>) -> tensor<16x8xf32> {104 %r = linalg.elementwise105 kind=#linalg.elementwise_kind<sub>106 indexing_maps = [affine_map<(d0, d1) -> (d1, d0)>,107 affine_map<(d0, d1) -> (d0, d1)>,108 affine_map<(d0, d1) -> (d0, d1)>]109 ins(%A, %B: tensor<8x16xf32>, tensor<16x8xf32>)110 outs(%C: tensor<16x8xf32>) -> tensor<16x8xf32>111 return %r : tensor<16x8xf32>112}113// -----114// CHECK-DAG: #[[IDENTITY:.+]] = affine_map<(d0, d1) -> (d0, d1)>115// CHECK-DAG: #[[TRANSPOSE:.+]] = affine_map<(d0, d1) -> (d1, d0)>116// CHECK-DAG: #[[BROADCAST:.+]] = affine_map<(d0, d1) -> (d0)>117//118// CHECK: @binary_transpose_a_broadcast_b(%[[A:.+]]: tensor<8x16xf32>, %[[B:.+]]: tensor<16xf32>, %[[C:.+]]: tensor<16x8xf32>)119// CHECK: linalg.generic120// CHECK-SAME: indexing_maps = [#[[TRANSPOSE]], #[[BROADCAST]], #[[IDENTITY]]]121// CHECK-SAME: iterator_types = ["parallel", "parallel"]122// CHECK-SAME: ins(%[[A]], %[[B]]123// CHECK-SAME: outs(%[[C]]124//125// CHECK: ^{{.*}}(%[[A_ARG:.+]]: f32, %[[B_ARG:.+]]: f32, %[[C_ARG:.+]]: f32)126// CHECK: %[[ADD:.+]] = arith.addf %[[A_ARG]], %[[B_ARG]] : f32127// CHECK: linalg.yield %[[ADD]] : f32128//129func.func @binary_transpose_a_broadcast_b(%A : tensor<8x16xf32>, %B: tensor<16xf32>, %C: tensor<16x8xf32>) -> tensor<16x8xf32> {130 %r = linalg.elementwise131 kind=#linalg.elementwise_kind<add>132 indexing_maps = [affine_map<(d0, d1) -> (d1, d0)>,133 affine_map<(d0, d1) -> (d0)>,134 affine_map<(d0, d1) -> (d0, d1)>]135 ins(%A, %B: tensor<8x16xf32>, tensor<16xf32>)136 outs(%C: tensor<16x8xf32>) -> tensor<16x8xf32>137 return %r : tensor<16x8xf32>138}139// -----140// CHECK-DAG: #[[IDENTITY:.+]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>141// CHECK-DAG: #[[PROJECTION:.+]] = affine_map<(d0, d1, d2) -> (d2, d1)>142//143// CHECK: @ternary(%[[A:.+]]: tensor<32x16xi1>, %[[B:.+]]: tensor<8x16x32xf32>, %[[C:.+]]: tensor<8x16x32xf32>, %[[D:.+]]: tensor<8x16x32xf32>)144// CHECK: linalg.generic145// CHECK-SAME: indexing_maps = [#[[PROJECTION]], #[[IDENTITY]], #[[IDENTITY]], #[[IDENTITY]]]146// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel"]147//148// CHECK-SAME: ins(%[[A]], %[[B]], %[[C]]149// CHECK-SAME: outs(%[[D]]150//151// CHECK: ^{{.*}}(%[[A_ARG:.+]]: i1, %[[B_ARG:.+]]: f32, %[[C_ARG:.+]]: f32, %[[D_ARG:.+]]: f32)152// CHECK: %[[SELECTED:.+]] = arith.select %[[A_ARG]], %[[B_ARG]], %[[C_ARG]] : f32153// CHECK: linalg.yield %[[SELECTED]] : f32154//155func.func @ternary(%A : tensor<32x16xi1>, %B: tensor<8x16x32xf32>, %C : tensor<8x16x32xf32>, %D : tensor<8x16x32xf32>) -> tensor<8x16x32xf32> {156 %r = linalg.elementwise157 kind=#linalg.elementwise_kind<select>158 indexing_maps = [affine_map<(d0, d1, d2) -> (d2, d1)>,159 affine_map<(d0, d1, d2) -> (d0, d1, d2)>,160 affine_map<(d0, d1, d2) -> (d0, d1, d2)>,161 affine_map<(d0, d1, d2) -> (d0, d1, d2)>]162 ins(%A, %B, %C : tensor<32x16xi1>, tensor<8x16x32xf32>, tensor<8x16x32xf32>)163 outs(%D: tensor<8x16x32xf32>) -> tensor<8x16x32xf32>164 return %r : tensor<8x16x32xf32>165}166