91 lines · plain
1// RUN: mlir-opt %s -split-input-file | FileCheck %s2//3// Note - the functions are named @{unary|binary}_{identity|transpose|broadcast|transpose_a|...}_{exp|mul|div|..}4 5// CHECK: @unary_identity_exp(%[[A:.+]]: tensor<8x16x32xf32>, %[[B:.+]]: tensor<8x16x32xf32>)6// CHECK: %{{.*}} = linalg.elementwise kind=#linalg.elementwise_kind<exp>7// CHECK-SAME ins(%[[A:.+]] : tensor<8x16x32xf32>) outs(%[[B:.+]] : tensor<8x16x32xf32>)8//9func.func @unary_identity_exp(%A : tensor<8x16x32xf32>, %B: tensor<8x16x32xf32>) -> tensor<8x16x32xf32> {10 %r = linalg.elementwise11 kind=#linalg.elementwise_kind<exp>12 ins(%A : tensor<8x16x32xf32>)13 outs(%B: tensor<8x16x32xf32>) -> tensor<8x16x32xf32>14 return %r : tensor<8x16x32xf32>15}16 17// -----18 19// CHECK-DAG: #[[IDENTITY:.+]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>20// CHECK-DAG: #[[PROJECTION:.+]] = affine_map<(d0, d1, d2) -> (d2, d1)>21//22// CHECK: @unary_projection_tanh(%[[A:.+]]: tensor<?x16xf32>,23// CHECK-SAME: %[[B:.+]]: tensor<8x16x?xf32>) -> tensor<8x16x?xf32> {24// CHECK: {{.*}} = linalg.elementwise kind=#linalg.elementwise_kind<tanh>25// CHECK-SAME: indexing_maps = [#[[PROJECTION]], #[[IDENTITY]]]26// CHECK-SAME: ins(%[[A]] : tensor<?x16xf32>) outs(%[[B]] : tensor<8x16x?xf32>) -> tensor<8x16x?xf32>27//28func.func @unary_projection_tanh(%A: tensor<?x16xf32>,29 %B: tensor<8x16x?xf32>) -> tensor<8x16x?xf32> {30 %r = linalg.elementwise31 kind=#linalg.elementwise_kind<tanh>32 indexing_maps = [affine_map<(d0, d1, d2) -> (d2, d1)>,33 affine_map<(d0, d1, d2) -> (d0, d1, d2)>]34 ins(%A : tensor<?x16xf32>)35 outs(%B: tensor<8x16x?xf32>) -> tensor<8x16x?xf32>36 return %r : tensor<8x16x?xf32>37}38 39// -----40 41// CHECK: @binary_identity_div(%[[A:.+]]: tensor<16x8xf32>, %[[B:.+]]: tensor<16x8xf32>,42// CHECK-SAME: %[[C:.+]]: tensor<16x8xf32>) -> tensor<16x8xf32> {43// CHECK: {{.*}} = linalg.elementwise44// CHECK-SAME: kind=#linalg.elementwise_kind<div>45// CHECK-SAME: ins(%[[A]], %[[B]] : tensor<16x8xf32>, tensor<16x8xf32>)46// CHECK-SAME: outs(%[[C]] : tensor<16x8xf32>) -> tensor<16x8xf32>47//48func.func @binary_identity_div(%A: tensor<16x8xf32>, %B: tensor<16x8xf32>,49 %C: tensor<16x8xf32>) -> tensor<16x8xf32> {50 %r = linalg.elementwise51 kind=#linalg.elementwise_kind<div>52 ins(%A, %B: tensor<16x8xf32>, tensor<16x8xf32>)53 outs(%C: tensor<16x8xf32>) -> tensor<16x8xf32>54 return %r : tensor<16x8xf32>55}56 57// -----58 59// CHECK: @binary_identity_mul_5Di(%[[A]]: tensor<1x2x3x4x5xi32>,60// CHECK-SAME: %[[B:.+]]: tensor<1x2x3x4x5xi32>,61// CHECK-SAME: %[[C:.+]]: tensor<1x2x3x4x5xi32>) -> tensor<1x2x3x4x5xi32> {62// CHECK: {{.*}} = linalg.elementwise63// CHECK-SAME: kind=#linalg.elementwise_kind<mul>64// CHECK-SAME: ins(%[[A]], %[[B]] : tensor<1x2x3x4x5xi32>, tensor<1x2x3x4x5xi32>)65// CHECK-SAME: outs(%[[C]] : tensor<1x2x3x4x5xi32>) -> tensor<1x2x3x4x5xi32>66//67func.func @binary_identity_mul_5Di(%A: tensor<1x2x3x4x5xi32>, %B: tensor<1x2x3x4x5xi32>,68 %C: tensor<1x2x3x4x5xi32>) -> tensor<1x2x3x4x5xi32> {69 %r = linalg.elementwise70 kind=#linalg.elementwise_kind<mul>71 ins(%A, %B: tensor<1x2x3x4x5xi32>, tensor<1x2x3x4x5xi32>)72 outs(%C: tensor<1x2x3x4x5xi32>) -> tensor<1x2x3x4x5xi32>73 return %r : tensor<1x2x3x4x5xi32>74}75 76// -----77 78// CHECK: @redundant_maps79// CHECK-NOT: indexing_maps80//81#map = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>82func.func @redundant_maps(%A: tensor<1x2x3x4x5xi32>, %B: tensor<1x2x3x4x5xi32>,83 %C: tensor<1x2x3x4x5xi32>) -> tensor<1x2x3x4x5xi32> {84 %r = linalg.elementwise85 kind=#linalg.elementwise_kind<mul>86 indexing_maps = [#map, #map, #map]87 ins(%A, %B: tensor<1x2x3x4x5xi32>, tensor<1x2x3x4x5xi32>)88 outs(%C: tensor<1x2x3x4x5xi32>) -> tensor<1x2x3x4x5xi32>89 return %r : tensor<1x2x3x4x5xi32>90}91