MLIR  20.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 "llvm/ADT/StringSet.h"
15 
16 #include "mlir/Dialect/Utils/DialectUtilsEnums.cpp.inc"
17 
18 using namespace mlir;
19 
20 bool mlir::isRowMajorMatmul(ArrayAttr indexingMaps) {
21  if (indexingMaps.size() != 3)
22  return false;
23 
24  AffineMap map0 = cast<AffineMapAttr>(indexingMaps[0]).getValue();
25  AffineMap map1 = cast<AffineMapAttr>(indexingMaps[1]).getValue();
26  AffineMap map2 = cast<AffineMapAttr>(indexingMaps[2]).getValue();
27 
28  if (map0.getNumResults() != 2 || map1.getNumResults() != 2 ||
29  map2.getNumResults() != 2 || map0.getNumInputs() != 3 ||
30  map1.getNumInputs() != 3 || map2.getNumInputs() != 3) {
31  return false;
32  }
33 
34  // Extract dimensions for MxK * KxN -> MxN
35  AffineExpr m = map2.getResult(0);
36  AffineExpr n = map2.getResult(1);
37  AffineExpr k = map0.getResult(1);
38  auto *context = indexingMaps.getContext();
39  auto mapA = AffineMapAttr::get(AffineMap::get(3, 0, {m, k}, context));
40  auto mapB = AffineMapAttr::get(AffineMap::get(3, 0, {k, n}, context));
41  auto mapC = AffineMapAttr::get(AffineMap::get(3, 0, {m, n}, context));
42  auto maps = ArrayAttr::get(context, {mapA, mapB, mapC});
43  return indexingMaps == maps;
44 }
45 
46 bool mlir::isColumnMajorMatmul(ArrayAttr indexingMaps) {
47  if (indexingMaps.size() != 3)
48  return false;
49 
50  AffineMap map0 = cast<AffineMapAttr>(indexingMaps[0]).getValue();
51  AffineMap map1 = cast<AffineMapAttr>(indexingMaps[1]).getValue();
52  AffineMap map2 = cast<AffineMapAttr>(indexingMaps[2]).getValue();
53 
54  if (map0.getNumResults() != 2 || map1.getNumResults() != 2 ||
55  map2.getNumResults() != 2 || map0.getNumInputs() != 3 ||
56  map1.getNumInputs() != 3 || map2.getNumInputs() != 3) {
57  return false;
58  }
59 
60  // Extract dimensions for KxM * NxK -> NxM
61  AffineExpr n = map2.getResult(0);
62  AffineExpr m = map2.getResult(1);
63  AffineExpr k = map0.getResult(0);
64  auto *context = indexingMaps.getContext();
65  auto mapA = AffineMapAttr::get(AffineMap::get(3, 0, {k, m}, context));
66  auto mapB = AffineMapAttr::get(AffineMap::get(3, 0, {n, k}, context));
67  auto mapC = AffineMapAttr::get(AffineMap::get(3, 0, {n, m}, context));
68  auto maps = ArrayAttr::get(context, {mapA, mapB, mapC});
69  return indexingMaps == maps;
70 }
71 
72 bool mlir::isRowMajorBatchMatmul(ArrayAttr indexingMaps) {
73  if (indexingMaps.size() != 3)
74  return false;
75 
76  AffineMap map0 = cast<AffineMapAttr>(indexingMaps[0]).getValue();
77  AffineMap map1 = cast<AffineMapAttr>(indexingMaps[1]).getValue();
78  AffineMap map2 = cast<AffineMapAttr>(indexingMaps[2]).getValue();
79 
80  if (map0.getNumResults() != 3 || map1.getNumResults() != 3 ||
81  map2.getNumResults() != 3 || map0.getNumInputs() != 4 ||
82  map1.getNumInputs() != 4 || map2.getNumInputs() != 4) {
83  return false;
84  }
85 
86  // Extract dimensions for BxMxK * BxKxN -> BxMxN
87  AffineExpr b = map2.getResult(0);
88  AffineExpr m = map2.getResult(1);
89  AffineExpr n = map2.getResult(2);
90  AffineExpr k = map0.getResult(2);
91  auto *context = indexingMaps.getContext();
92  auto mapA = AffineMapAttr::get(AffineMap::get(4, 0, {b, m, k}, context));
93  auto mapB = AffineMapAttr::get(AffineMap::get(4, 0, {b, k, n}, context));
94  auto mapC = AffineMapAttr::get(AffineMap::get(4, 0, {b, m, n}, context));
95  auto maps = ArrayAttr::get(context, {mapA, mapB, mapC});
96  return indexingMaps == maps;
97 }
98 
99 bool mlir::isVecmat(ArrayAttr indexingMaps) {
100  if (indexingMaps.size() != 3)
101  return false;
102  AffineMap map0 = cast<AffineMapAttr>(indexingMaps[0]).getValue();
103  AffineMap map1 = cast<AffineMapAttr>(indexingMaps[1]).getValue();
104  AffineMap map2 = cast<AffineMapAttr>(indexingMaps[2]).getValue();
105 
106  if (map0.getNumResults() != 1 || map1.getNumResults() != 2 ||
107  map2.getNumResults() != 1 || map0.getNumInputs() != 2 ||
108  map1.getNumInputs() != 2 || map2.getNumInputs() != 2) {
109  return false;
110  }
111 
112  // Extract dimensions for K * KxN -> N
113  AffineExpr k = map0.getResult(0);
114  AffineExpr n = map2.getResult(0);
115  auto *context = indexingMaps.getContext();
116  auto mapA = AffineMapAttr::get(AffineMap::get(2, 0, {k}, context));
117  auto mapB = AffineMapAttr::get(AffineMap::get(2, 0, {k, n}, context));
118  auto mapC = AffineMapAttr::get(AffineMap::get(2, 0, {n}, context));
119  auto maps = ArrayAttr::get(context, {mapA, mapB, mapC});
120  return indexingMaps == maps;
121 }
122 
123 bool mlir::isBatchVecmat(ArrayAttr indexingMaps) {
124  if (indexingMaps.size() != 3)
125  return false;
126  AffineMap map0 = cast<AffineMapAttr>(indexingMaps[0]).getValue();
127  AffineMap map1 = cast<AffineMapAttr>(indexingMaps[1]).getValue();
128  AffineMap map2 = cast<AffineMapAttr>(indexingMaps[2]).getValue();
129 
130  if (map0.getNumResults() != 2 || map1.getNumResults() != 3 ||
131  map2.getNumResults() != 2 || map0.getNumInputs() != 3 ||
132  map1.getNumInputs() != 3 || map2.getNumInputs() != 3) {
133  return false;
134  }
135 
136  // Extract dimensions for B*K * B*K*N -> B*N
137  AffineExpr b = map0.getResult(0);
138  AffineExpr k = map0.getResult(1);
139  AffineExpr n = map2.getResult(1);
140  auto *context = indexingMaps.getContext();
141  auto mapA = AffineMapAttr::get(AffineMap::get(3, 0, {b, k}, context));
142  auto mapB = AffineMapAttr::get(AffineMap::get(3, 0, {b, k, n}, context));
143  auto mapC = AffineMapAttr::get(AffineMap::get(3, 0, {b, n}, context));
144  auto maps = ArrayAttr::get(context, {mapA, mapB, mapC});
145  return indexingMaps == maps;
146 }
147 
148 bool mlir::isMatvec(ArrayAttr indexingMaps) {
149  if (indexingMaps.size() != 3)
150  return false;
151  AffineMap map0 = cast<AffineMapAttr>(indexingMaps[0]).getValue();
152  AffineMap map1 = cast<AffineMapAttr>(indexingMaps[1]).getValue();
153  AffineMap map2 = cast<AffineMapAttr>(indexingMaps[2]).getValue();
154 
155  if (map0.getNumResults() != 2 || map1.getNumResults() != 1 ||
156  map2.getNumResults() != 1 || map0.getNumInputs() != 2 ||
157  map1.getNumInputs() != 2 || map2.getNumInputs() != 2) {
158  return false;
159  }
160 
161  // Extract dimensions for N*K * K -> N
162  AffineExpr k = map1.getResult(0);
163  AffineExpr n = map2.getResult(0);
164  auto *context = indexingMaps.getContext();
165  auto mapA = AffineMapAttr::get(AffineMap::get(2, 0, {n, k}, context));
166  auto mapB = AffineMapAttr::get(AffineMap::get(2, 0, {k}, context));
167  auto mapC = AffineMapAttr::get(AffineMap::get(2, 0, {n}, context));
168  auto maps = ArrayAttr::get(context, {mapA, mapB, mapC});
169  return indexingMaps == maps;
170 }
171 
172 bool mlir::isBatchMatvec(ArrayAttr indexingMaps) {
173  if (indexingMaps.size() != 3)
174  return false;
175  AffineMap map0 = cast<AffineMapAttr>(indexingMaps[0]).getValue();
176  AffineMap map1 = cast<AffineMapAttr>(indexingMaps[1]).getValue();
177  AffineMap map2 = cast<AffineMapAttr>(indexingMaps[2]).getValue();
178 
179  if (map0.getNumResults() != 3 || map1.getNumResults() != 2 ||
180  map2.getNumResults() != 2 || map0.getNumInputs() != 3 ||
181  map1.getNumInputs() != 3 || map2.getNumInputs() != 3) {
182  return false;
183  }
184 
185  // Extract dimensions for B*N*K * B*K -> B*N
186  AffineExpr b = map0.getResult(0);
187  AffineExpr k = map1.getResult(1);
188  AffineExpr n = map2.getResult(1);
189  auto *context = indexingMaps.getContext();
190  auto mapA = AffineMapAttr::get(AffineMap::get(3, 0, {b, n, k}, context));
191  auto mapB = AffineMapAttr::get(AffineMap::get(3, 0, {b, k}, context));
192  auto mapC = AffineMapAttr::get(AffineMap::get(3, 0, {b, n}, context));
193  auto maps = ArrayAttr::get(context, {mapA, mapB, mapC});
194  return indexingMaps == maps;
195 }
196 
197 Operation *mlir::clone(OpBuilder &b, Operation *op, TypeRange newResultTypes,
198  ValueRange newOperands) {
199  IRMapping bvm;
200  OperationState state(op->getLoc(), op->getName(), newOperands, newResultTypes,
201  op->getAttrs());
202  for (Region &r : op->getRegions()) {
203  Region *newRegion = state.addRegion();
204  b.cloneRegionBefore(r, *newRegion, newRegion->begin(), bvm);
205  }
206  return b.create(state);
207 }
208 
209 Operation *mlir::cloneWithoutRegions(OpBuilder &b, Operation *op,
210  TypeRange newResultTypes,
211  ValueRange newOperands) {
212  OperationState state(op->getLoc(), op->getName(), newOperands, newResultTypes,
213  op->getAttrs());
214  for (size_t cnt = 0, e = op->getNumRegions(); cnt < e; ++cnt)
215  state.addRegion();
216  return b.create(state);
217 }
218 
219 SmallVector<NamedAttribute>
220 mlir::getPrunedAttributeList(Operation *op, ArrayRef<StringRef> elidedAttrs) {
221  llvm::StringSet<> elidedAttrsSet;
222  elidedAttrsSet.insert(elidedAttrs.begin(), elidedAttrs.end());
223  SmallVector<NamedAttribute> attrs;
224  for (auto attr : op->getAttrs()) {
225  if (elidedAttrsSet.count(attr.getName()))
226  continue;
227  attrs.push_back(attr);
228  }
229  return attrs;
230 }
mlir::AffineExpr
Base type for affine expression.
Definition: AffineExpr.h:68
mlir::AffineExpr::getContext
MLIRContext * getContext() const
Definition: AffineExpr.cpp:33
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:1263
mlir::AffineMap::getNumResults
unsigned getNumResults() const
Definition: AffineMap.cpp:402
mlir::AffineMap::getNumInputs
unsigned getNumInputs() const
Definition: AffineMap.cpp:403
mlir::AffineMap::getResult
AffineExpr getResult(unsigned idx) const
Definition: AffineMap.cpp:411
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:215
mlir::OpBuilder::cloneRegionBefore
void cloneRegionBefore(Region &region, Region &parent, Region::iterator before, IRMapping &mapping)
Clone the blocks that belong to "region" before the given position in another region "parent".
Definition: Builders.cpp:615
mlir::OpBuilder::create
Operation * create(const OperationState &state)
Creates an operation given the fields represented as an OperationState.
Definition: Builders.cpp:497
mlir::Operation
Operation is the basic unit of execution within MLIR.
Definition: Operation.h:88
mlir::Operation::getNumRegions
unsigned getNumRegions()
Returns the number of regions held by this operation.
Definition: Operation.h:669
mlir::Operation::getLoc
Location getLoc()
The source location the operation was defined or derived from.
Definition: Operation.h:223
mlir::Operation::getAttrs
ArrayRef< NamedAttribute > getAttrs()
Return all of the attributes on this operation.
Definition: Operation.h:507
mlir::Operation::getRegions
MutableArrayRef< Region > getRegions()
Returns the regions held by this operation.
Definition: Operation.h:672
mlir::Operation::getName
OperationName getName()
The name of an operation is the key identifier for it.
Definition: Operation.h:119
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:36
mlir::ValueRange
This class provides an abstraction over the different types of ranges over Values.
Definition: ValueRange.h:381
mlir
Include the generated interface declarations.
Definition: LocalAliasAnalysis.h:20
mlir::isColumnMajorMatmul
bool isColumnMajorMatmul(ArrayAttr indexingMaps)
Tests whether the given maps describe a column major matmul.
Definition: StructuredOpsUtils.cpp:46
mlir::isBatchMatvec
bool isBatchMatvec(ArrayAttr indexingMaps)
Tests whether the given maps describe a batch matrix vector multiplication.
Definition: StructuredOpsUtils.cpp:172
mlir::cloneWithoutRegions
Operation * cloneWithoutRegions(OpBuilder &b, Operation *op, TypeRange newResultTypes, ValueRange newOperands)
Definition: StructuredOpsUtils.cpp:209
mlir::isMatvec
bool isMatvec(ArrayAttr indexingMaps)
Tests whether the given maps describe a matrix vector multiplication.
Definition: StructuredOpsUtils.cpp:148
mlir::isBatchVecmat
bool isBatchVecmat(ArrayAttr indexingMaps)
Tests whether the given maps describe a batch vector matrix multiplication.
Definition: StructuredOpsUtils.cpp:123
mlir::isRowMajorBatchMatmul
bool isRowMajorBatchMatmul(ArrayAttr indexingMaps)
Tests whether the given maps describe a row major batch matmul.
Definition: StructuredOpsUtils.cpp:72
mlir::clone
Operation * clone(OpBuilder &b, Operation *op, TypeRange newResultTypes, ValueRange newOperands)
Definition: StructuredOpsUtils.cpp:197
mlir::get
auto get(MLIRContext *context, Ts &&...params)
Helper method that injects context only if needed, this helps unify some of the attribute constructio...
Definition: BytecodeImplementation.h:509
mlir::isVecmat
bool isVecmat(ArrayAttr indexingMaps)
Tests whether the given maps describe a vector matrix multiplication.
Definition: StructuredOpsUtils.cpp:99
mlir::isRowMajorMatmul
bool isRowMajorMatmul(ArrayAttr indexingMaps)
Tests whether the given maps describe a row major matmul.
Definition: StructuredOpsUtils.cpp:20
mlir::getPrunedAttributeList
SmallVector< NamedAttribute > getPrunedAttributeList(Operation *op, ArrayRef< StringRef > elidedAttrs)
Definition: StructuredOpsUtils.cpp:220
mlir::OperationState
This represents an operation in an abstracted form, suitable for use with the builder APIs.
Definition: OperationSupport.h:950
Generated on Tue Nov 19 2024 04:32:30 for MLIR by doxygen 1.9.1