52 lines · plain
1// RUN: mlir-opt %s -linalg-block-pack-matmul="block-factors=32,16,64 \2// RUN: lhs-transpose-outer-blocks=false lhs-transpose-inner-blocks=false \3// RUN: rhs-transpose-outer-blocks=true rhs-transpose-inner-blocks=true" \4// RUN: -canonicalize | FileCheck %s --check-prefix=MMT4D5 6// RUN: mlir-opt %s -linalg-block-pack-matmul="block-factors=32,16,64 \7// RUN: lhs-transpose-outer-blocks=false lhs-transpose-inner-blocks=false \8// RUN: rhs-transpose-outer-blocks=false rhs-transpose-inner-blocks=false" \9// RUN: -canonicalize | FileCheck %s --check-prefix=MM4D10 11// RUN: mlir-opt %s -linalg-block-pack-matmul="block-factors=32,16,64 \12// RUN: lhs-transpose-outer-blocks=true lhs-transpose-inner-blocks=true \13// RUN: rhs-transpose-outer-blocks=false rhs-transpose-inner-blocks=false" \14// RUN: -canonicalize | FileCheck %s --check-prefix=MTM4D15 16func.func @block_matmul(17 %A: tensor<64x128xf32>, %B: tensor<128x64xf32>, %C: tensor<64x64xf32>) -> tensor<64x64xf32> {18 %0 = linalg.matmul ins(%A, %B : tensor<64x128xf32>, tensor<128x64xf32>)19 outs(%C : tensor<64x64xf32>) -> tensor<64x64xf32>20 return %0 : tensor<64x64xf32>21}22 23// MMT4D-DAG: #[[$MAP:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d2, d3, d5)>24// MMT4D-DAG: #[[$MAP1:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d1, d2, d4, d5)>25// MMT4D-DAG: #[[$MAP2:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1, d3, d4)>26// MMT4D-LABEL: func @block_matmul27// MMT4D-COUNT-3: linalg.pack28// MMT4D: linalg.generic29// MMT4D-SAME: indexing_maps = [#[[$MAP]], #[[$MAP1]], #[[$MAP2]]]30// MMT4D-SAME: iterator_types = ["parallel", "parallel", "reduction", "parallel", "parallel", "reduction"]31// MMT4D-COUNT-1: linalg.unpack32 33// MM4D-DAG: #[[$MAP:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d2, d3, d5)>34// MM4D-DAG: #[[$MAP1:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d2, d1, d5, d4)>35// MM4D-DAG: #[[$MAP2:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1, d3, d4)>36// MM4D-LABEL: func @block_matmul37// MM4D-COUNT-3: linalg.pack38// MM4D: linalg.generic39// MM4D-SAME: indexing_maps = [#[[$MAP]], #[[$MAP1]], #[[$MAP2]]]40// MM4D-SAME: iterator_types = ["parallel", "parallel", "reduction", "parallel", "parallel", "reduction"]41// MM4D-COUNT-1: linalg.unpack42 43// MTM4D-DAG: #[[$MAP:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d2, d0, d5, d3)>44// MTM4D-DAG: #[[$MAP1:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d2, d1, d5, d4)>45// MTM4D-DAG: #[[$MAP2:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1, d3, d4)>46// MTM4D-LABEL: func @block_matmul47// MTM4D-COUNT-3: linalg.pack48// MTM4D: linalg.generic49// MTM4D-SAME: indexing_maps = [#[[$MAP]], #[[$MAP1]], #[[$MAP2]]]50// MTM4D-SAME: iterator_types = ["parallel", "parallel", "reduction", "parallel", "parallel", "reduction"]51// MTM4D-COUNT-1: linalg.unpack52