54 lines · plain
1// RUN: mlir-opt %s -transform-interpreter | FileCheck %s2 3func.func @matmul_to_outerproduct(%A: memref<3x4xf32>, %B: memref<4x3xf32>, %C: memref<3x3xf32>) {4 linalg.matmul ins(%A, %B: memref<3x4xf32>, memref<4x3xf32>)5 outs(%C: memref<3x3xf32>)6 return7}8 9// CHECK-LABEL: func.func @matmul_to_outerproduct(10// CHECK-SAME: %[[A:.*]]: memref<3x4xf32>,11// CHECK-SAME: %[[B:.*]]: memref<4x3xf32>,12// CHECK-SAME: %[[C:.*]]: memref<3x3xf32>) {13// CHECK: %[[VEC_A:.*]] = vector.transfer_read %[[A]]14// CHECK: %[[VEC_B:.*]] = vector.transfer_read %[[B]]15// CHECK: %[[VEC_C:.*]] = vector.transfer_read %[[C]]16// CHECK: %[[VEC_A_T:.*]] = vector.transpose %[[VEC_A]], [1, 0] : vector<3x4xf32> to vector<4x3xf32>17// CHECK: %[[A0:.*]] = vector.extract %[[VEC_A_T]][0] : vector<3xf32> from vector<4x3xf32>18// CHECK: %[[B0:.*]] = vector.extract %[[VEC_B]][0] : vector<3xf32> from vector<4x3xf32>19// CHECK: %[[OP_0:.*]] = vector.outerproduct %[[A0]], %[[B0]], %[[VEC_C]]20// CHECK: %[[A1:.*]] = vector.extract %[[VEC_A_T]][1] : vector<3xf32> from vector<4x3xf32>21// CHECK: %[[B1:.*]] = vector.extract %[[VEC_B]][1] : vector<3xf32> from vector<4x3xf32>22// CHECK: %[[OP_1:.*]] = vector.outerproduct %[[A1]], %[[B1]], %[[OP_0]]23// CHECK: %[[A_2:.*]] = vector.extract %[[VEC_A_T]][2] : vector<3xf32> from vector<4x3xf32>24// CHECK: %[[B_2:.*]] = vector.extract %[[VEC_B]][2] : vector<3xf32> from vector<4x3xf32>25// CHECK: %[[OP_2:.*]] = vector.outerproduct %[[A_2]], %[[B_2]], %[[OP_1]]26// CHECK: %[[A_3:.*]] = vector.extract %[[VEC_A_T]][3] : vector<3xf32> from vector<4x3xf32>27// CHECK: %[[B_3:.*]] = vector.extract %[[VEC_B]][3] : vector<3xf32> from vector<4x3xf32>28// CHECK: %[[RES:.*]] = vector.outerproduct %[[A_3]], %[[B_3]], %[[OP_2]]29// CHECK: vector.transfer_write %[[RES]], %[[C]]{{.*}} : vector<3x3xf32>, memref<3x3xf32>30 31module attributes {transform.with_named_sequence} {32 transform.named_sequence @__transform_main(%module: !transform.any_op {transform.readonly}) {33 %func = transform.structured.match ops{["func.func"]} in %module : (!transform.any_op) -> !transform.any_op34 35 // Vectorize: linalg.matmul -> vector.multi_reduction36 %matmul = transform.structured.match ops{["linalg.matmul"]} in %func : (!transform.any_op) -> !transform.any_op37 transform.structured.vectorize %matmul : !transform.any_op38 39 // vector.multi_reduction --> vector.contract40 transform.apply_patterns to %func {41 transform.apply_patterns.vector.reduction_to_contract42 // Reduce the rank of xfer ops. This transform vector.contract to be more43 // more matmul-like and to enable the lowering to outer product Ops.44 transform.apply_patterns.vector.transfer_permutation_patterns45 } : !transform.any_op46 47 // vector.contract --> vector.outerproduct48 transform.apply_patterns to %func {49 transform.apply_patterns.vector.lower_contraction lowering_strategy = "outerproduct"50 } : !transform.any_op51 transform.yield52 }53}54