44 lines · plain
1// RUN: mlir-opt %s -linalg-fold-into-elementwise -split-input-file | FileCheck %s2 3// CHECK-DAG: #[[IDENTITY:.+]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>4// CHECK-DAG: #[[TRANSPOSED:.+]] = affine_map<(d0, d1, d2) -> (d1, d0, d2)>5//6// CHECK: func.func @unary_transpose(%[[A:.+]]: tensor<16x8x32xf32>, %[[B:.+]]: tensor<8x16x32xf32>) -> tensor<8x16x32xf32> {7// CHECK-NEXT: %[[RES:.+]] = linalg.elementwise kind=#linalg.elementwise_kind<exp>8// CHECK-SAME: indexing_maps = [#[[TRANSPOSED]], #[[IDENTITY]]]9// CHECK-SAME: ins(%[[A]] : tensor<16x8x32xf32>) outs(%[[B]] : tensor<8x16x32xf32>) -> tensor<8x16x32xf32>10// CHECK-NEXT: return %[[RES]] : tensor<8x16x32xf32>11//12func.func @unary_transpose(%A : tensor<16x8x32xf32>, %B: tensor<8x16x32xf32>) -> tensor<8x16x32xf32> {13 %empty = tensor.empty() : tensor<8x16x32xf32>14 %transposed_A = linalg.transpose ins(%A : tensor<16x8x32xf32>) outs(%empty : tensor<8x16x32xf32>) permutation = [1, 0, 2]15 %result = linalg.elementwise kind=#linalg.elementwise_kind<exp>16 ins(%transposed_A : tensor<8x16x32xf32>) outs(%B: tensor<8x16x32xf32>) -> tensor<8x16x32xf32>17 return %result : tensor<8x16x32xf32>18}19 20// -----21 22// CHECK-DAG: #[[IDENTITY:.+]] = affine_map<(d0, d1) -> (d0, d1)>23// CHECK-DAG: #[[TRANSPOSED:.+]] = affine_map<(d0, d1) -> (d1, d0)>24//25// CHECK: func.func @binary_transposed(%[[A:.+]]: tensor<?x?xf32>, %[[B:.+]]: tensor<?x?xf32>, %[[C:.+]]: tensor<?x?xf32>) -> tensor<?x?xf32> {26// CHECK-NEXT: %[[RES:.+]] = linalg.elementwise kind=#linalg.elementwise_kind<add>27// CHECK-SAME: indexing_maps = [#[[IDENTITY]], #[[TRANSPOSED]], #[[IDENTITY]]]28// CHECK-SAME: ins(%[[A]], %[[B]] : tensor<?x?xf32>, tensor<?x?xf32>) outs(%[[C]] : tensor<?x?xf32>) -> tensor<?x?xf32>29// CHECK-NEXT: return %[[RES]] : tensor<?x?xf32>30//31func.func @binary_transposed(%A : tensor<?x?xf32>, %B: tensor<?x?xf32>, %C: tensor<?x?xf32>) -> tensor<?x?xf32> {32 %c0 = arith.constant 0 : index33 %c1 = arith.constant 1 : index34 %dim0 = tensor.dim %A, %c0 : tensor<?x?xf32>35 %dim1 = tensor.dim %A, %c1 : tensor<?x?xf32>36 37 %empty = tensor.empty(%dim1, %dim0) : tensor<?x?xf32>38 %transposed_B = linalg.transpose ins(%B : tensor<?x?xf32>) outs(%empty : tensor<?x?xf32>) permutation = [1, 0]39 %result = linalg.elementwise kind=#linalg.elementwise_kind<add>40 ins(%A, %transposed_B : tensor<?x?xf32>, tensor<?x?xf32>)41 outs(%C: tensor<?x?xf32>) -> tensor<?x?xf32>42 return %result : tensor<?x?xf32>43}44