MLIR 24.0.0git
StructuredOpsUtils.cpp
Go to the documentation of this file.
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/Builders.h"
12#include "mlir/IR/BuiltinAttributes.h"
13#include "mlir/IR/IRMapping.h"
14#include "mlir/IR/OpImplementation.h"
15#include "llvm/ADT/StringSet.h"
16
17#include "mlir/Dialect/Utils/DialectUtilsEnums.cpp.inc"
18
19using namespace mlir;
20
21bool mlir::isRowMajorMatmul(ArrayAttr indexingMaps) {
22 if (indexingMaps.size() != 3)
23 return false;
24
25 AffineMap map0 = cast<AffineMapAttr>(indexingMaps[0]).getValue();
26 AffineMap map1 = cast<AffineMapAttr>(indexingMaps[1]).getValue();
27 AffineMap map2 = cast<AffineMapAttr>(indexingMaps[2]).getValue();
28
29 if (map0.getNumResults() != 2 || map1.getNumResults() != 2 ||
30 map2.getNumResults() != 2 || map0.getNumInputs() != 3 ||
31 map1.getNumInputs() != 3 || map2.getNumInputs() != 3) {
32 return false;
33 }
34
35 // Extract dimensions for MxK * KxN -> MxN
36 AffineExpr m = map2.getResult(0);
37 AffineExpr n = map2.getResult(1);
38 AffineExpr k = map0.getResult(1);
39 auto *context = indexingMaps.getContext();
40 auto mapA = AffineMapAttr::get(AffineMap::get(3, 0, {m, k}, context));
41 auto mapB = AffineMapAttr::get(AffineMap::get(3, 0, {k, n}, context));
42 auto mapC = AffineMapAttr::get(AffineMap::get(3, 0, {m, n}, context));
43 auto maps = ArrayAttr::get(context, {mapA, mapB, mapC});
44 return indexingMaps == maps;
45}
46
47bool mlir::isColumnMajorMatmul(ArrayAttr indexingMaps) {
48 if (indexingMaps.size() != 3)
49 return false;
50
51 AffineMap map0 = cast<AffineMapAttr>(indexingMaps[0]).getValue();
52 AffineMap map1 = cast<AffineMapAttr>(indexingMaps[1]).getValue();
53 AffineMap map2 = cast<AffineMapAttr>(indexingMaps[2]).getValue();
54
55 if (map0.getNumResults() != 2 || map1.getNumResults() != 2 ||
56 map2.getNumResults() != 2 || map0.getNumInputs() != 3 ||
57 map1.getNumInputs() != 3 || map2.getNumInputs() != 3) {
58 return false;
59 }
60
61 // Extract dimensions for KxM * NxK -> NxM
62 AffineExpr n = map2.getResult(0);
63 AffineExpr m = map2.getResult(1);
64 AffineExpr k = map0.getResult(0);
65 auto *context = indexingMaps.getContext();
66 auto mapA = AffineMapAttr::get(AffineMap::get(3, 0, {k, m}, context));
67 auto mapB = AffineMapAttr::get(AffineMap::get(3, 0, {n, k}, context));
68 auto mapC = AffineMapAttr::get(AffineMap::get(3, 0, {n, m}, context));
69 auto maps = ArrayAttr::get(context, {mapA, mapB, mapC});
70 return indexingMaps == maps;
71}
72
73bool mlir::isRowMajorBatchMatmul(ArrayAttr indexingMaps) {
74 if (indexingMaps.size() != 3)
75 return false;
76
77 AffineMap map0 = cast<AffineMapAttr>(indexingMaps[0]).getValue();
78 AffineMap map1 = cast<AffineMapAttr>(indexingMaps[1]).getValue();
79 AffineMap map2 = cast<AffineMapAttr>(indexingMaps[2]).getValue();
80
81 if (map0.getNumResults() != 3 || map1.getNumResults() != 3 ||
82 map2.getNumResults() != 3 || map0.getNumInputs() != 4 ||
83 map1.getNumInputs() != 4 || map2.getNumInputs() != 4) {
84 return false;
85 }
86
87 // Extract dimensions for BxMxK * BxKxN -> BxMxN
88 AffineExpr b = map2.getResult(0);
89 AffineExpr m = map2.getResult(1);
90 AffineExpr n = map2.getResult(2);
91 AffineExpr k = map0.getResult(2);
92 auto *context = indexingMaps.getContext();
93 auto mapA = AffineMapAttr::get(AffineMap::get(4, 0, {b, m, k}, context));
94 auto mapB = AffineMapAttr::get(AffineMap::get(4, 0, {b, k, n}, context));
95 auto mapC = AffineMapAttr::get(AffineMap::get(4, 0, {b, m, n}, context));
96 auto maps = ArrayAttr::get(context, {mapA, mapB, mapC});
97 return indexingMaps == maps;
98}
99
100bool mlir::isVecmat(ArrayAttr indexingMaps) {
101 if (indexingMaps.size() != 3)
102 return false;
103 AffineMap map0 = cast<AffineMapAttr>(indexingMaps[0]).getValue();
104 AffineMap map1 = cast<AffineMapAttr>(indexingMaps[1]).getValue();
105 AffineMap map2 = cast<AffineMapAttr>(indexingMaps[2]).getValue();
106
107 if (map0.getNumResults() != 1 || map1.getNumResults() != 2 ||
108 map2.getNumResults() != 1 || map0.getNumInputs() != 2 ||
109 map1.getNumInputs() != 2 || map2.getNumInputs() != 2) {
110 return false;
111 }
112
113 // Extract dimensions for K * KxN -> N
114 AffineExpr k = map0.getResult(0);
115 AffineExpr n = map2.getResult(0);
116 auto *context = indexingMaps.getContext();
117 auto mapA = AffineMapAttr::get(AffineMap::get(2, 0, {k}, context));
118 auto mapB = AffineMapAttr::get(AffineMap::get(2, 0, {k, n}, context));
119 auto mapC = AffineMapAttr::get(AffineMap::get(2, 0, {n}, context));
120 auto maps = ArrayAttr::get(context, {mapA, mapB, mapC});
121 return indexingMaps == maps;
122}
123
124bool mlir::isBatchVecmat(ArrayAttr indexingMaps) {
125 if (indexingMaps.size() != 3)
126 return false;
127 AffineMap map0 = cast<AffineMapAttr>(indexingMaps[0]).getValue();
128 AffineMap map1 = cast<AffineMapAttr>(indexingMaps[1]).getValue();
129 AffineMap map2 = cast<AffineMapAttr>(indexingMaps[2]).getValue();
130
131 if (map0.getNumResults() != 2 || map1.getNumResults() != 3 ||
132 map2.getNumResults() != 2 || map0.getNumInputs() != 3 ||
133 map1.getNumInputs() != 3 || map2.getNumInputs() != 3) {
134 return false;
135 }
136
137 // Extract dimensions for B*K * B*K*N -> B*N
138 AffineExpr b = map0.getResult(0);
139 AffineExpr k = map0.getResult(1);
140 AffineExpr n = map2.getResult(1);
141 auto *context = indexingMaps.getContext();
142 auto mapA = AffineMapAttr::get(AffineMap::get(3, 0, {b, k}, context));
143 auto mapB = AffineMapAttr::get(AffineMap::get(3, 0, {b, k, n}, context));
144 auto mapC = AffineMapAttr::get(AffineMap::get(3, 0, {b, n}, context));
145 auto maps = ArrayAttr::get(context, {mapA, mapB, mapC});
146 return indexingMaps == maps;
147}
148
149bool mlir::isMatvec(ArrayAttr indexingMaps) {
150 if (indexingMaps.size() != 3)
151 return false;
152 AffineMap map0 = cast<AffineMapAttr>(indexingMaps[0]).getValue();
153 AffineMap map1 = cast<AffineMapAttr>(indexingMaps[1]).getValue();
154 AffineMap map2 = cast<AffineMapAttr>(indexingMaps[2]).getValue();
155
156 if (map0.getNumResults() != 2 || map1.getNumResults() != 1 ||
157 map2.getNumResults() != 1 || map0.getNumInputs() != 2 ||
158 map1.getNumInputs() != 2 || map2.getNumInputs() != 2) {
159 return false;
160 }
161
162 // Extract dimensions for N*K * K -> N
163 AffineExpr k = map1.getResult(0);
164 AffineExpr n = map2.getResult(0);
165 auto *context = indexingMaps.getContext();
166 auto mapA = AffineMapAttr::get(AffineMap::get(2, 0, {n, k}, context));
167 auto mapB = AffineMapAttr::get(AffineMap::get(2, 0, {k}, context));
168 auto mapC = AffineMapAttr::get(AffineMap::get(2, 0, {n}, context));
169 auto maps = ArrayAttr::get(context, {mapA, mapB, mapC});
170 return indexingMaps == maps;
171}
172
173bool mlir::isBatchMatvec(ArrayAttr indexingMaps) {
174 if (indexingMaps.size() != 3)
175 return false;
176 AffineMap map0 = cast<AffineMapAttr>(indexingMaps[0]).getValue();
177 AffineMap map1 = cast<AffineMapAttr>(indexingMaps[1]).getValue();
178 AffineMap map2 = cast<AffineMapAttr>(indexingMaps[2]).getValue();
179
180 if (map0.getNumResults() != 3 || map1.getNumResults() != 2 ||
181 map2.getNumResults() != 2 || map0.getNumInputs() != 3 ||
182 map1.getNumInputs() != 3 || map2.getNumInputs() != 3) {
183 return false;
184 }
185
186 // Extract dimensions for B*N*K * B*K -> B*N
187 AffineExpr b = map0.getResult(0);
188 AffineExpr k = map1.getResult(1);
189 AffineExpr n = map2.getResult(1);
190 auto *context = indexingMaps.getContext();
191 auto mapA = AffineMapAttr::get(AffineMap::get(3, 0, {b, n, k}, context));
192 auto mapB = AffineMapAttr::get(AffineMap::get(3, 0, {b, k}, context));
193 auto mapC = AffineMapAttr::get(AffineMap::get(3, 0, {b, n}, context));
194 auto maps = ArrayAttr::get(context, {mapA, mapB, mapC});
195 return indexingMaps == maps;
196}
197
198Operation *mlir::clone(OpBuilder &b, Operation *op, TypeRange newResultTypes,
199 ValueRange newOperands) {
200 IRMapping bvm;
201 OperationState state(op->getLoc(), op->getName(), newOperands, newResultTypes,
202 op->getDiscardableAttrDictionary().getValue());
203 state.propertiesAttr = op->getPropertiesAsAttribute();
204 for (Region &r : op->getRegions()) {
205 Region *newRegion = state.addRegion();
206 b.cloneRegionBefore(r, *newRegion, newRegion->begin(), bvm);
207 }
208 return b.create(state);
209}
210
211Operation *mlir::cloneWithoutRegions(OpBuilder &b, Operation *op,
212 TypeRange newResultTypes,
213 ValueRange newOperands) {
214 OperationState state(op->getLoc(), op->getName(), newOperands, newResultTypes,
215 op->getDiscardableAttrDictionary().getValue());
216 state.propertiesAttr = op->getPropertiesAsAttribute();
217 for (size_t cnt = 0, e = op->getNumRegions(); cnt < e; ++cnt)
218 state.addRegion();
219 return b.create(state);
220}
221
222SmallVector<NamedAttribute>
223mlir::getPrunedAttributeList(Operation *op, ArrayRef<StringRef> elidedAttrs) {
224 llvm::StringSet<> elidedAttrsSet;
225 elidedAttrsSet.insert_range(elidedAttrs);
226 SmallVector<NamedAttribute> attrs;
227 for (auto attr : op->getDiscardableAttrDictionary().getValue()) {
228 if (elidedAttrsSet.count(attr.getName()))
229 continue;
230 attrs.push_back(attr);
231 }
232 return attrs;
233}
b
b
Return true if permutation is a valid permutation of the outer_dims_perm (case OuterOrInnerPerm::Oute...
Definition LinalgTransformOps.cpp:2154
mlir::AffineExpr
Base type for affine expression.
Definition AffineExpr.h:68
mlir::AffineMap
A multi-dimensional affine map Affine map's are immutable like Type's, and they are uniqued.
Definition AffineMap.h:46
mlir::AffineMap::get
static AffineMap get(MLIRContext *context)
Returns a zero result affine map with no dimensions or symbols: () -> ().
Definition MLIRContext.cpp:1334
mlir::AffineMap::getNumResults
unsigned getNumResults() const
Definition AffineMap.cpp:400
mlir::AffineMap::getNumInputs
unsigned getNumInputs() const
Definition AffineMap.cpp:401
mlir::AffineMap::getResult
AffineExpr getResult(unsigned idx) const
Definition AffineMap.cpp:409
mlir::IRMapping
This is a utility class for mapping one set of IR entities to another.
Definition IRMapping.h:26
mlir::OpBuilder
This class helps build Operations.
Definition Builders.h:210
mlir::Operation
Operation is the basic unit of execution within MLIR.
Definition Operation.h:87
mlir::Operation::getNumRegions
unsigned getNumRegions()
Returns the number of regions held by this operation.
Definition Operation.h:726
mlir::Operation::getLoc
Location getLoc()
The source location the operation was defined or derived from.
Definition Operation.h:240
mlir::Operation::getPropertiesAsAttribute
Attribute getPropertiesAsAttribute()
Return the properties converted to an attribute.
Definition Operation.cpp:345
mlir::Operation::getName
OperationName getName()
The name of an operation is the key identifier for it.
Definition Operation.h:115
mlir::Operation::getDiscardableAttrDictionary
DictionaryAttr getDiscardableAttrDictionary()
Return all of the discardable attributes on this operation as a DictionaryAttr.
Definition Operation.h:553
mlir::Operation::getRegions
MutableArrayRef< Region > getRegions()
Returns the regions held by this operation.
Definition Operation.h:729
mlir::Region
This class contains a list of basic blocks and a link to the parent operation it is attached to.
Definition Region.h:26
mlir::Region::begin
iterator begin()
Definition Region.h:55
mlir::TypeRange
This class provides an abstraction over the various different ranges of value types.
Definition TypeRange.h:40
mlir::ValueRange
This class provides an abstraction over the different types of ranges over Values.
Definition ValueRange.h:389
mlir
Include the generated interface declarations.
Definition ABIRewriteContext.h:32
mlir::isColumnMajorMatmul
bool isColumnMajorMatmul(ArrayAttr indexingMaps)
Tests whether the given maps describe a column major matmul.
Definition StructuredOpsUtils.cpp:47
mlir::isBatchMatvec
bool isBatchMatvec(ArrayAttr indexingMaps)
Tests whether the given maps describe a batch matrix vector multiplication.
Definition StructuredOpsUtils.cpp:173
mlir::cloneWithoutRegions
Operation * cloneWithoutRegions(OpBuilder &b, Operation *op, TypeRange newResultTypes, ValueRange newOperands)
Definition StructuredOpsUtils.cpp:211
mlir::isMatvec
bool isMatvec(ArrayAttr indexingMaps)
Tests whether the given maps describe a matrix vector multiplication.
Definition StructuredOpsUtils.cpp:149
mlir::isBatchVecmat
bool isBatchVecmat(ArrayAttr indexingMaps)
Tests whether the given maps describe a batch vector matrix multiplication.
Definition StructuredOpsUtils.cpp:124
mlir::isRowMajorBatchMatmul
bool isRowMajorBatchMatmul(ArrayAttr indexingMaps)
Tests whether the given maps describe a row major batch matmul.
Definition StructuredOpsUtils.cpp:73
mlir::clone
Operation * clone(OpBuilder &b, Operation *op, TypeRange newResultTypes, ValueRange newOperands)
Definition StructuredOpsUtils.cpp:198
mlir::isVecmat
bool isVecmat(ArrayAttr indexingMaps)
Tests whether the given maps describe a vector matrix multiplication.
Definition StructuredOpsUtils.cpp:100
mlir::isRowMajorMatmul
bool isRowMajorMatmul(ArrayAttr indexingMaps)
Tests whether the given maps describe a row major matmul.
Definition StructuredOpsUtils.cpp:21
mlir::getPrunedAttributeList
SmallVector< NamedAttribute > getPrunedAttributeList(Operation *op, ArrayRef< StringRef > elidedAttrs)
Definition StructuredOpsUtils.cpp:223
mlir::OperationState
This represents an operation in an abstracted form, suitable for use with the builder APIs.
Definition OperationSupport.h:1020
mlir::OperationState::propertiesAttr
Attribute propertiesAttr
This Attribute is used to opaquely construct the properties of the operation.
Definition OperationSupport.h:1038
mlir::OperationState::addRegion
Region * addRegion()
Create a region that should be attached to the operation.
Definition OperationSupport.cpp:267
Generated on for MLIR by doxygen 1.14.0