1 //===- StructuredOpsUtils.cpp - Utilities used by structured ops ----------===//
2 //
3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4 // See https://llvm.org/LICENSE.txt for license information.
5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6 //
7 //===----------------------------------------------------------------------===//
8
9 #include "mlir/Dialect/Utils/StructuredOpsUtils.h"
10 #include "mlir/IR/AffineMap.h"
11 #include "mlir/IR/BuiltinAttributes.h"
12
13 using namespace mlir;
14
isRowMajorMatmul(ArrayAttr indexingMaps)15 bool mlir::isRowMajorMatmul(ArrayAttr indexingMaps) {
16 if (indexingMaps.size() != 3)
17 return false;
18
19 auto map0 = indexingMaps[0].cast<AffineMapAttr>().getValue();
20 auto map1 = indexingMaps[1].cast<AffineMapAttr>().getValue();
21 auto map2 = indexingMaps[2].cast<AffineMapAttr>().getValue();
22
23 if (map0.getNumResults() != 2 || map1.getNumResults() != 2 ||
24 map2.getNumResults() != 2 || map0.getNumInputs() != 3 ||
25 map1.getNumInputs() != 3 || map2.getNumInputs() != 3) {
26 return false;
27 }
28
29 // Extract dimensions for MxK * KxN -> MxN
30 AffineExpr m = map2.getResult(0);
31 AffineExpr n = map2.getResult(1);
32 AffineExpr k = map0.getResult(1);
33 auto *context = indexingMaps.getContext();
34 auto mapA = AffineMapAttr::get(AffineMap::get(3, 0, {m, k}, context));
35 auto mapB = AffineMapAttr::get(AffineMap::get(3, 0, {k, n}, context));
36 auto mapC = AffineMapAttr::get(AffineMap::get(3, 0, {m, n}, context));
37 auto maps = ArrayAttr::get(context, {mapA, mapB, mapC});
38 return indexingMaps == maps;
39 }
40
isColumnMajorMatmul(ArrayAttr indexingMaps)41 bool mlir::isColumnMajorMatmul(ArrayAttr indexingMaps) {
42 if (indexingMaps.size() != 3)
43 return false;
44
45 auto map0 = indexingMaps[0].cast<AffineMapAttr>().getValue();
46 auto map1 = indexingMaps[1].cast<AffineMapAttr>().getValue();
47 auto map2 = indexingMaps[2].cast<AffineMapAttr>().getValue();
48
49 if (map0.getNumResults() != 2 || map1.getNumResults() != 2 ||
50 map2.getNumResults() != 2 || map0.getNumInputs() != 3 ||
51 map1.getNumInputs() != 3 || map2.getNumInputs() != 3) {
52 return false;
53 }
54
55 // Extract dimensions for KxM * NxK -> NxM
56 AffineExpr n = map2.getResult(0);
57 AffineExpr m = map2.getResult(1);
58 AffineExpr k = map0.getResult(0);
59 auto *context = indexingMaps.getContext();
60 auto mapA = AffineMapAttr::get(AffineMap::get(3, 0, {k, m}, context));
61 auto mapB = AffineMapAttr::get(AffineMap::get(3, 0, {n, k}, context));
62 auto mapC = AffineMapAttr::get(AffineMap::get(3, 0, {n, m}, context));
63 auto maps = ArrayAttr::get(context, {mapA, mapB, mapC});
64 return indexingMaps == maps;
65 }
66
isRowMajorBatchMatmul(ArrayAttr indexingMaps)67 bool mlir::isRowMajorBatchMatmul(ArrayAttr indexingMaps) {
68 if (indexingMaps.size() != 3)
69 return false;
70
71 auto map0 = indexingMaps[0].cast<AffineMapAttr>().getValue();
72 auto map1 = indexingMaps[1].cast<AffineMapAttr>().getValue();
73 auto map2 = indexingMaps[2].cast<AffineMapAttr>().getValue();
74
75 if (map0.getNumResults() != 3 || map1.getNumResults() != 3 ||
76 map2.getNumResults() != 3 || map0.getNumInputs() != 4 ||
77 map1.getNumInputs() != 4 || map2.getNumInputs() != 4) {
78 return false;
79 }
80
81 // Extract dimensions for BxMxK * BxKxN -> BxMxN
82 AffineExpr b = map2.getResult(0);
83 AffineExpr m = map2.getResult(1);
84 AffineExpr n = map2.getResult(2);
85 AffineExpr k = map0.getResult(2);
86 auto *context = indexingMaps.getContext();
87 auto mapA = AffineMapAttr::get(AffineMap::get(4, 0, {b, m, k}, context));
88 auto mapB = AffineMapAttr::get(AffineMap::get(4, 0, {b, k, n}, context));
89 auto mapC = AffineMapAttr::get(AffineMap::get(4, 0, {b, m, n}, context));
90 auto maps = ArrayAttr::get(context, {mapA, mapB, mapC});
91 return indexingMaps == maps;
92 }
93