MLIR 23.0.0git
Specialize.cpp
Go to the documentation of this file.
1//===- Specialize.cpp - linalg generic ops to named 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// This file implements a method to specialize generic operations to named
10// operations. Conceptually it is the opposite of generalize.cpp.
11//
12//===----------------------------------------------------------------------===//
13
23
24namespace mlir {
25#define GEN_PASS_DEF_LINALGSPECIALIZEGENERICOPSPASS
26#include "mlir/Dialect/Linalg/Passes.h.inc"
27} // namespace mlir
28
29#define DEBUG_TYPE "linalg-specialization"
30
31#define REPLACE_BINARY_OP(NEWOP, OPERANDS_SWAP) \
32 (rewriter.replaceOpWithNewOp<NEWOP>( \
33 genericOp, \
34 ValueRange{genericOp.getDpsInputs()[(OPERANDS_SWAP) ? 1 : 0], \
35 genericOp.getDpsInputs()[(OPERANDS_SWAP) ? 0 : 1]}, \
36 ValueRange{genericOp.getDpsInits()[0]}))
37
38#define REPLACE_UNARY_OP(NEWOP) \
39 (rewriter.replaceOpWithNewOp<NEWOP>(genericOp, \
40 ValueRange{genericOp.getDpsInputs()[0]}, \
41 ValueRange{genericOp.getDpsInits()[0]}))
42
43using namespace mlir;
44using namespace mlir::linalg;
45
46// Given a elementwise single binary linalg generic op, checks whether the
47// binary op accesses operands as swapped. e.g.
48// this differentiates between a linalg-generic body that contains:
49// ^bb0(%a: f32, %b: f32, %c : f32):
50// %0 = arith.subf %a, %b : f32
51// linalg.yield %0: f32
52// against:
53// ^bb0(%a: f32, %b: f32, %c : f32):
54// %0 = arith.subf %b, %a : f32
55// linalg.yield %0: f32
56// Former is linalg.sub(a,b), latter is linalg.sub(b,a).
57static bool areBinOpsSwapped(GenericOp genericOp) {
58 Block *body = genericOp.getBody();
59 Operation *op = &body->front();
60 bool swapped = false;
61 if (op->getOpOperand(0).get() != body->getArgument(0)) {
62 swapped = true;
63 assert(op->getOpOperand(0).get() == body->getArgument(1) &&
64 op->getOpOperand(1).get() == body->getArgument(0) &&
65 "binary op uses just one block arg");
66 }
67 return swapped;
68}
69
70//===----------------------------------------------------------------------===//
71// Specialize linalg generic to matmul variants.
72//===----------------------------------------------------------------------===//
73/// Identifies linalg.generic that is essentially named op of the form:
74// ` linalg.{batch_}?matmul{_transpose_a | _transpose_b}? `
75//
76// It is possible that a linalg.generic may be implementing a matmul but not
77// in a straight-forward way e.g. below is matrix multiply over some slice
78// ```
79// %0 = linalg.generic {
80// indexing_maps = [affine_map<(d0, d1, d2) -> (3, d1, d0)>,
81// affine_map<(d0, d1, d2) -> (d0, 5, d2)>,
82// affine_map<(d0, d1, d2) -> (d2, d1, 13)>],
83// iterator_types = ["parallel", "parallel", "parallel"]}
84// ins(%A, %B : tensor<20x20x20xf32>, tensor<20x20x20xf32>)
85// outs(%C : tensor<20x20x20xf32>) {
86// ^bb0(%a: f32, %b: f32, %c : f32):
87// %mul = arith.mulf %a, %b : f32
88// %add = arith.addf %mul, %c : f32
89// linalg.yield %add : f32
90// } -> tensor<20x20x20xf32>
91// ```
92// It is not possible to represent above as named op.
93// e.g. linalg.batch_matmul(%A, %B : tensor<20x20x20xf32>, ...) is
94// not the same as linalg.generic above.
95namespace {
96enum class IndexMatchResult {
97 Match = 0, // identity map.
98 Transposed, // transposed map.
99 Mismatch // none of the above.
100};
101
102// Checks whether the input Affine `map` contains two consecutive dims that
103// can be interpreted as accessing a 2D matrix. It is assumed that the row
104// column dimension are adjacent axis (in this order) and start at
105// `rowDimIdx` in the input map.
106//
107// e.g. consider A matrix in `C[M,N] = A[M,K] * B[K,N]`. We will check
108// whether the map of A is identity (match), transposed, or something
109// completely different (mis-match). Similar for B and C.
110static IndexMatchResult matchOperandMap(AffineMap map, unsigned rowDimIdx,
111 unsigned expectedPosOfRowDim,
112 unsigned expectedPosOfColDim) {
113 // Get the matrix multiply indices. They are past the batch indices.
114 auto exprOfRowDim = map.getResults()[rowDimIdx];
115 auto exprOfColDim = map.getResults()[rowDimIdx + 1];
116
117 // They should be pure dimension ids.
118 if (exprOfRowDim.getKind() != AffineExprKind::DimId ||
119 exprOfColDim.getKind() != AffineExprKind::DimId)
120 return IndexMatchResult::Mismatch;
121
122 auto posRowDim = cast<AffineDimExpr>(exprOfRowDim).getPosition();
123 auto posColDim = cast<AffineDimExpr>(exprOfColDim).getPosition();
124
125 if (expectedPosOfRowDim == posRowDim && expectedPosOfColDim == posColDim)
126 return IndexMatchResult::Match;
127
128 if (expectedPosOfRowDim == posColDim && expectedPosOfColDim == posRowDim)
129 return IndexMatchResult::Transposed;
130
131 return IndexMatchResult::Mismatch;
132}
133
134// Replaces genericOp with `NamedOpTy` op, supplied as a template arg.
135// All the variants expressed as pseudo regular expression:
136// `linalg.{batch_}?matmul{_transpose_a | _transpose_b}?`
137// have same number of ins/out, so its easy to stamp different versions.
138// `castTy` is an optional type function that indicates whether (and which) cast
139// attribute is needed for the named matmul op variant.
140template <typename NamedOpTy>
141static LinalgOp replaceWithMatmulVariant(RewriterBase &rewriter, GenericOp op,
142 std::optional<TypeFn> castTy) {
143 SmallVector<NamedAttribute> castAttrVec;
144 // Only explicitly specify the cast attribute for unsigned cast; signed is
145 // the default for linalg.matmul/linalg.batch_matmul.
146 if (castTy.has_value() && *castTy == TypeFn::cast_unsigned)
147 castAttrVec = {rewriter.getNamedAttr(
148 "cast", TypeFnAttr::get(rewriter.getContext(), *castTy))};
149
150 LinalgOp namedOp = rewriter.replaceOpWithNewOp<NamedOpTy>(
151 op, ValueRange{op.getDpsInputs()[0], op.getDpsInputs()[1]},
152 ValueRange{op.getDpsInits()[0]}, castAttrVec);
153 return namedOp;
154}
155
156// Returns the cast type to use for a matmul-like named op. If the generic
157// contains casts that cannot be represented (e.g. output casts or mixed
158// signedness), return std::nullopt.
159static std::optional<TypeFn> getCastTypeForMatmulLikeOp(GenericOp genericOp) {
160 bool foundCastForMatmulOutput = false;
161 SmallVector<TypeFn> castTyFns;
162 genericOp.getBody()->walk([&](CastOpInterface castOp) {
163 // Collect forward slice of the cast op to check if it is for the matmul
164 // output.
165 SetVector<Operation *> forwardSlice;
166 getForwardSlice(castOp, &forwardSlice);
167
168 // If there is no multiplication op in the forward slice, then this cast
169 // op is for the matmul output. Cast ops on matmul output cannot be
170 // expressed by the matmul op variant.
171 if (!llvm::any_of(forwardSlice, [](Operation *op) {
172 // We check explicitly for these multiplication ops in
173 // `specializeLinalgContractions()` to infer matmul-like ops.
174 return isa<arith::MulIOp, arith::MulFOp, complex::MulOp>(op);
175 })) {
176 foundCastForMatmulOutput = true;
177 return WalkResult::interrupt();
178 }
179
180 // Determine the cast type.
181 if (isa<arith::ExtUIOp, arith::UIToFPOp, arith::FPToUIOp>(castOp))
182 castTyFns.push_back(TypeFn::cast_unsigned);
183 else if (isa<arith::ExtSIOp, arith::SIToFPOp, arith::FPToSIOp>(castOp))
184 castTyFns.push_back(TypeFn::cast_signed);
185
186 return WalkResult::advance();
187 });
188
189 if (foundCastForMatmulOutput)
190 return std::nullopt;
191
192 if (!castTyFns.empty()) {
193 // If there were multiple different cast types found, then we can't express
194 // them using matmul-like ops. They only allow a single cast type for all
195 // inputs.
196 if (!llvm::all_equal(castTyFns))
197 return std::nullopt;
198 return castTyFns.front();
199 }
200
201 // Default to signed cast for matmul-like ops.
202 return TypeFn::cast_signed;
203}
204
205// Converts linalg.generic to named linalg.*matmul* where possible.
206static FailureOr<LinalgOp> specializeLinalgContractions(RewriterBase &rewriter,
207 GenericOp genericOp) {
208 if (genericOp.getNumDpsInputs() != 2 || genericOp.getNumDpsInits() != 1)
209 return failure();
210
211 // Early exit if not projected permutations.
212 auto mapRange = genericOp.getIndexingMapsArray();
213 if (llvm::any_of(mapRange,
214 [](AffineMap m) { return !m.isProjectedPermutation(); }))
215 return failure();
216
217 // Linalg generic contraction can be across multiple axis e.g.
218 // ```
219 // linalg.generic
220 // {indexing_maps = [affine_map<(m, n, k1, k2) -> (m, k1, k2)>,
221 // affine_map<(m, n, k1, k2) -> (k2, k1, n)>,
222 // affine_map<(m, n, k1, k2) -> (m, n)>],
223 // iterator_types = ["parallel", "parallel",
224 // "reduction", "reduction"]}
225 // ins(%A, %B : tensor<10x20x30xf32>, tensor<30x20x40xf32>)
226 // outs(%C : tensor<10x40xf32>) {
227 // ^bb0(%a: f32, %b: f32, %c: f32):
228 // %1 = arith.mulf %a, %b : f32
229 // %2 = arith.addf %c, %1 : f32
230 // linalg.yield %2 : f32
231 // } -> tensor<10x40xf32>
232 // ```
233 // In above contraction, there are two reduction dimensions {k1, k2}
234 // and although a valid linalg contraction, it is not a named-op
235 // matrix multiply kind. Therefore, reject multi-dim reduction.
236 auto res = inferContractionDims(genericOp);
237 if (!succeeded(res))
238 return failure();
239 auto dims = *res;
240 if (dims.m.size() != 1 || dims.n.size() != 1 || dims.k.size() != 1)
241 return failure();
242
244 *genericOp.getBlock(), [](Operation *first, Operation *second) {
245 return (isa<arith::MulFOp>(first) && isa<arith::AddFOp>(second)) ||
246 (isa<arith::MulIOp>(first) && isa<arith::AddIOp>(second)) ||
247 (isa<complex::MulOp>(first) && isa<complex::AddOp>(second));
248 }))
249 return failure();
250
251 // Check rank of operands
252 auto indexingMaps = genericOp.getIndexingMapsArray();
253 if (llvm::any_of(indexingMaps, [&dims](AffineMap m) {
254 return m.getResults().size() !=
255 dims.batch.size() + 2 /* any two of {m,n,k} */;
256 }))
257 return failure();
258
259 auto numOfBatchDims = dims.batch.size();
260 if (indexingMaps[0].getNumDims() != numOfBatchDims + 3)
261 return failure();
262
263 if (numOfBatchDims) {
264 // Each operand in a linalg generic contraction could express different
265 // permutations for its batch dimension. But for named op it must be
266 // identity since separate maps are not specified.
267 if (llvm::any_of(indexingMaps, [numOfBatchDims](AffineMap m) {
268 for (unsigned i = 0; i < numOfBatchDims; ++i) {
269 auto expr = m.getResults()[i];
270 if (expr.getKind() != AffineExprKind::DimId ||
271 cast<AffineDimExpr>(expr).getPosition() != i)
272 return true;
273 }
274 return false;
275 }))
276 return failure();
277 }
278
279 auto a =
280 matchOperandMap(indexingMaps[0], numOfBatchDims, dims.m[0], dims.k[0]);
281 auto b =
282 matchOperandMap(indexingMaps[1], numOfBatchDims, dims.k[0], dims.n[0]);
283 auto c =
284 matchOperandMap(indexingMaps[2], numOfBatchDims, dims.m[0], dims.n[0]);
285
286 if (llvm::is_contained({a, b, c}, IndexMatchResult::Mismatch))
287 return failure();
288
289 if (c != IndexMatchResult::Match ||
290 (a == IndexMatchResult::Transposed && b == IndexMatchResult::Transposed))
291 return failure();
292
293 // Determine the cast type for the named matmul op, or bail out if casts
294 // cannot be represented by the named op.
295 std::optional<TypeFn> castTy = getCastTypeForMatmulLikeOp(genericOp);
296 if (!castTy)
297 return rewriter.notifyMatchFailure(
298 genericOp, "contains invalid cast ops for the named matmul op");
299
300 /// Codegen the different matmul variants.
301 if (numOfBatchDims) {
302 return replaceWithMatmulVariant<BatchMatmulOp>(rewriter, genericOp, castTy);
303 }
304 return replaceWithMatmulVariant<MatmulOp>(rewriter, genericOp, castTy);
305}
306
307/// Utility to specialize a `genericOp` with a convolution op of type `ConvOpTy`
308/// with `dilations` and `strides`.
309template <typename ConvOpTy>
310static FailureOr<LinalgOp>
311specializeToConvOp(RewriterBase &rewriter, GenericOp genericOp,
312 ArrayRef<int64_t> dilations, ArrayRef<int64_t> strides) {
313 SmallVector<Value> inputs = genericOp.getDpsInputs();
314 ValueRange outputs = genericOp.getDpsInits();
315 SmallVector<Type> resultTypes = genericOp.hasPureTensorSemantics()
316 ? TypeRange(ValueRange(outputs))
317 : TypeRange{};
318 LinalgOp namedOp;
319 // Ops with no dilations and no strides.
320 if constexpr (std::is_same_v<ConvOpTy, linalg::Conv1DOp> ||
321 std::is_same_v<ConvOpTy, linalg::Conv2DOp> ||
322 std::is_same_v<ConvOpTy, linalg::Conv3DOp>) {
323 namedOp = rewriter.replaceOpWithNewOp<ConvOpTy>(genericOp, resultTypes,
324 inputs, outputs);
325 } else {
326 Attribute stridesAttr = rewriter.getI64TensorAttr(strides);
327 Attribute dilationsAttr = rewriter.getI64TensorAttr(dilations);
328 namedOp = rewriter.replaceOpWithNewOp<ConvOpTy>(
329 genericOp, resultTypes, inputs, outputs, stridesAttr, dilationsAttr);
330 }
331 return namedOp;
332}
333
334/// Converts linalg.generic to named linalg.*conv/pooling* where possible.
335static FailureOr<LinalgOp> specializeLinalgConvolutions(RewriterBase &rewriter,
336 GenericOp genericOp) {
337#define CONV_OP_SPECIALIZER(ConvOpTy) \
338 if (std::optional<DilationsAndStrides> convParams = \
339 matchConvolutionOpOfType<ConvOpTy>(genericOp)) \
340 return specializeToConvOp<ConvOpTy>( \
341 rewriter, genericOp, convParams->dilations, convParams->strides); \
342 // -----------------------------
343 // Convolution ops.
344 // -----------------------------
345 CONV_OP_SPECIALIZER(linalg::Conv1DOp);
346 CONV_OP_SPECIALIZER(linalg::Conv1DNwcWcfOp);
347 CONV_OP_SPECIALIZER(linalg::Conv1DNcwFcwOp);
348 CONV_OP_SPECIALIZER(linalg::Conv2DOp);
349 CONV_OP_SPECIALIZER(linalg::Conv2DNhwcHwcfOp);
350 CONV_OP_SPECIALIZER(linalg::Conv2DNhwcHwcfQOp);
351 CONV_OP_SPECIALIZER(linalg::Conv2DNhwcFhwcOp);
352 CONV_OP_SPECIALIZER(linalg::Conv2DNhwcFhwcQOp);
353 CONV_OP_SPECIALIZER(linalg::Conv2DNchwFchwOp);
354 CONV_OP_SPECIALIZER(linalg::Conv2DNchwFchwQOp);
355 CONV_OP_SPECIALIZER(linalg::Conv2DNgchwFgchwOp);
356 CONV_OP_SPECIALIZER(linalg::Conv2DNgchwGfchwOp);
357 CONV_OP_SPECIALIZER(linalg::Conv2DNgchwGfchwQOp);
358 CONV_OP_SPECIALIZER(linalg::Conv2DNhwgcGfhwcOp);
359 CONV_OP_SPECIALIZER(linalg::Conv2DNhwgcGfhwcQOp);
360 CONV_OP_SPECIALIZER(linalg::Conv3DOp);
361 CONV_OP_SPECIALIZER(linalg::Conv3DNdhwcDhwcfOp);
362 CONV_OP_SPECIALIZER(linalg::Conv3DNdhwcDhwcfQOp);
363 CONV_OP_SPECIALIZER(linalg::Conv3DNcdhwFcdhwOp);
364 // -----------------------------
365 // Depthwise Convolution ops.
366 // -----------------------------
367 CONV_OP_SPECIALIZER(linalg::DepthwiseConv1DNcwCwOp);
368 CONV_OP_SPECIALIZER(linalg::DepthwiseConv1DNwcWcOp);
369 CONV_OP_SPECIALIZER(linalg::DepthwiseConv1DNwcWcmOp);
370 CONV_OP_SPECIALIZER(linalg::DepthwiseConv2DNchwChwOp);
371 CONV_OP_SPECIALIZER(linalg::DepthwiseConv2DNhwcHwcOp);
372 CONV_OP_SPECIALIZER(linalg::DepthwiseConv2DNhwcHwcQOp);
373 CONV_OP_SPECIALIZER(linalg::DepthwiseConv2DNhwcHwcmOp);
374 CONV_OP_SPECIALIZER(linalg::DepthwiseConv2DNhwcHwcmQOp);
375 CONV_OP_SPECIALIZER(linalg::DepthwiseConv3DNdhwcDhwcOp);
376 CONV_OP_SPECIALIZER(linalg::DepthwiseConv3DNcdhwCdhwOp);
377 CONV_OP_SPECIALIZER(linalg::DepthwiseConv3DNdhwcDhwcmOp);
378 // -----------------------------
379 // Pooling ops.
380 // -----------------------------
381 CONV_OP_SPECIALIZER(linalg::PoolingNhwcMaxOp);
382 CONV_OP_SPECIALIZER(linalg::PoolingNhwcMinOp);
383 CONV_OP_SPECIALIZER(linalg::PoolingNhwcSumOp);
384 CONV_OP_SPECIALIZER(linalg::PoolingNhwcMaxUnsignedOp);
385 CONV_OP_SPECIALIZER(linalg::PoolingNhwcMinUnsignedOp);
386 CONV_OP_SPECIALIZER(linalg::PoolingNchwSumOp);
387 CONV_OP_SPECIALIZER(linalg::PoolingNchwMaxOp);
388 CONV_OP_SPECIALIZER(linalg::PoolingNwcSumOp);
389 CONV_OP_SPECIALIZER(linalg::PoolingNcwSumOp);
390 CONV_OP_SPECIALIZER(linalg::PoolingNwcMaxOp);
391 CONV_OP_SPECIALIZER(linalg::PoolingNwcMaxUnsignedOp);
392 CONV_OP_SPECIALIZER(linalg::PoolingNcwMaxOp);
393 CONV_OP_SPECIALIZER(linalg::PoolingNwcMinOp);
394 CONV_OP_SPECIALIZER(linalg::PoolingNwcMinUnsignedOp);
395 CONV_OP_SPECIALIZER(linalg::PoolingNdhwcSumOp);
396 CONV_OP_SPECIALIZER(linalg::PoolingNdhwcMaxOp);
397 CONV_OP_SPECIALIZER(linalg::PoolingNdhwcMinOp);
398#undef CONV_OP_SPECIALIZER
399 return failure();
400}
401
402} // namespace
403
404//===----------------------------------------------------------------------===//
405// Categorize linalg generic to named op where possible.
406//===----------------------------------------------------------------------===//
408 GenericOp genericOp) {
409 // Copy
410 if (isaCopyOpInterface(genericOp)) {
411 LinalgOp namedOp = rewriter.replaceOpWithNewOp<CopyOp>(
412 genericOp, genericOp.getDpsInputs()[0], genericOp.getDpsInits()[0]);
413 return namedOp;
414 }
415
416 // Fill
417 if (std::optional<Value> fillValue = isaFillOpInterface(genericOp)) {
418 // Always use the detected fill value, regardless of pattern
419 LinalgOp namedOp = rewriter.replaceOpWithNewOp<FillOp>(
420 genericOp, *fillValue, genericOp.getDpsInits()[0]);
421 return namedOp;
422 }
423
424 // Broadcast
425 std::optional<SmallVector<int64_t>> equivalentToBroadcast =
426 isaBroadcastOpInterface(genericOp);
427 if (equivalentToBroadcast) {
428 auto dims = *equivalentToBroadcast;
429 LinalgOp namedOp = rewriter.replaceOpWithNewOp<BroadcastOp>(
430 genericOp, genericOp.getDpsInputs()[0], genericOp.getDpsInits()[0],
431 dims);
432 return namedOp;
433 }
434
435 // Transpose
436 std::optional<SmallVector<int64_t>> equivalentToTranspose =
437 isaTransposeOpInterface(genericOp);
438 if (equivalentToTranspose) {
439 auto permutation = *equivalentToTranspose;
440 LinalgOp namedOp = rewriter.replaceOpWithNewOp<TransposeOp>(
441 genericOp, genericOp.getDpsInputs()[0], genericOp.getDpsInits()[0],
442 permutation);
443 return namedOp;
444 }
445
446 // Elementwise Unary
447 if (isaElemwiseSingleUnaryOpInterface(genericOp)) {
448 Operation *op = &genericOp.getBody()->front();
449 if (isa<math::ExpOp>(op)) {
450 LinalgOp namedOp = REPLACE_UNARY_OP(ExpOp);
451 return namedOp;
452 }
453 }
454
455 // Elementwise Binary
456 if (isaElemwiseSingleBinaryOpInterface(genericOp)) {
457 bool swap = areBinOpsSwapped(genericOp);
458 Operation *op = &genericOp.getBody()->front();
459 if (isa<arith::AddFOp>(op)) {
460 LinalgOp namedOp = REPLACE_BINARY_OP(AddOp, swap);
461 return namedOp;
462 }
463 if (isa<arith::SubFOp>(op)) {
464 LinalgOp namedOp = REPLACE_BINARY_OP(SubOp, swap);
465 return namedOp;
466 }
467 if (isa<arith::MulFOp>(op)) {
468 LinalgOp namedOp = REPLACE_BINARY_OP(MulOp, swap);
469 return namedOp;
470 }
471 if (isa<arith::DivFOp>(op)) {
472 LinalgOp namedOp = REPLACE_BINARY_OP(DivOp, swap);
473 return namedOp;
474 }
475 }
476
477 // Contraction - e.g. matmul
478 if (isaContractionOpInterface(genericOp)) {
479 return specializeLinalgContractions(rewriter, genericOp);
480 }
481
482 // Convolution - e.g. *conv/pooling*
483 if (isaConvolutionOpInterface(genericOp)) {
484 return specializeLinalgConvolutions(rewriter, genericOp);
485 }
486 return failure();
487}
488
489namespace {
490struct LinalgSpecializeGenericOpsPass
491 : public impl::LinalgSpecializeGenericOpsPassBase<
492 LinalgSpecializeGenericOpsPass> {
493
494 using impl::LinalgSpecializeGenericOpsPassBase<
495 LinalgSpecializeGenericOpsPass>::LinalgSpecializeGenericOpsPassBase;
496 void runOnOperation() override;
497};
498} // namespace
499
500void LinalgSpecializeGenericOpsPass::runOnOperation() {
501 RewritePatternSet patterns(&getContext());
504
505 if (failed(applyPatternsGreedily(getOperation(), std::move(patterns))))
506 signalPassFailure();
507}
508
b
Return true if permutation is a valid permutation of the outer_dims_perm (case OuterOrInnerPerm::Oute...
b getContext())
#define REPLACE_BINARY_OP(NEWOP, OPERANDS_SWAP)
#define CONV_OP_SPECIALIZER(ConvOpTy)
static bool areBinOpsSwapped(GenericOp genericOp)
#define REPLACE_UNARY_OP(NEWOP)
A multi-dimensional affine map Affine map's are immutable like Type's, and they are uniqued.
Definition AffineMap.h:46
bool isProjectedPermutation(bool allowZeroInResults=false) const
Returns true if the AffineMap represents a subset (i.e.
ArrayRef< AffineExpr > getResults() const
Attributes are known-constant values of operations.
Definition Attributes.h:25
Block represents an ordered list of Operations.
Definition Block.h:33
BlockArgument getArgument(unsigned i)
Definition Block.h:139
Operation & front()
Definition Block.h:163
DenseIntElementsAttr getI64TensorAttr(ArrayRef< int64_t > values)
Definition Builders.cpp:186
MLIRContext * getContext() const
Definition Builders.h:56
NamedAttribute getNamedAttr(StringRef name, Attribute val)
Definition Builders.cpp:94
IRValueT get() const
Return the current value being used by this operand.
Operation is the basic unit of execution within MLIR.
Definition Operation.h:88
OpOperand & getOpOperand(unsigned idx)
Definition Operation.h:388
This class coordinates the application of a rewrite on a set of IR, providing a way for clients to tr...
std::enable_if_t<!std::is_convertible< CallbackT, Twine >::value, LogicalResult > notifyMatchFailure(Location loc, CallbackT &&reasonCallback)
Used to notify the listener that the IR failed to be rewritten because of a match failure,...
OpTy replaceOpWithNewOp(Operation *op, Args &&...args)
Replace the results of the given (original) op with a new op that is created without verification (re...
This class provides an abstraction over the various different ranges of value types.
Definition TypeRange.h:37
This class provides an abstraction over the different types of ranges over Values.
Definition ValueRange.h:387
static WalkResult advance()
Definition WalkResult.h:47
static WalkResult interrupt()
Definition WalkResult.h:46
bool isContractionBody(Block &block, function_ref< bool(Operation *, Operation *)> isaPair, llvm::raw_ostream &errs=mlir::thread_safe_nulls())
Returns true if the block contains a contraction of the following form:
std::optional< SmallVector< int64_t > > isaTransposeOpInterface(GenericOp genericOp)
Checks whether genericOp is semantically equivalent to a linalg.transpose.
bool isaElemwiseSingleUnaryOpInterface(GenericOp genericOp)
Checks whether a given genericOp is semantically equivalent to a single linalgelementwise unary op.
bool isaCopyOpInterface(LinalgOp linalgOp)
Checks whether linalgOp is semantically equivalent to a linalg.copyOp.
void populateDecomposeProjectedPermutationPatterns(RewritePatternSet &patterns)
Add patterns to make explicit broadcasts and transforms in the input operands of a genericOp.
FailureOr< LinalgOp > specializeGenericOp(RewriterBase &rewriter, GenericOp genericOp)
Create a namedOp from the given GenericOp and replace the GenericOp.
bool isaConvolutionOpInterface(LinalgOp linalgOp, bool allowEmptyConvolvedDims=false)
Checks whether linalgOp conforms to ConvolutionOpInterface.
std::optional< SmallVector< int64_t > > isaBroadcastOpInterface(GenericOp genericOp)
Checks whether genericOp is semantically equivalent to a linalg.broadcast.
FailureOr< ContractionDimensions > inferContractionDims(LinalgOp linalgOp)
Find at least 2 parallel (m and n) and 1 reduction (k) dimension candidates that form a matmul subcom...
bool isaContractionOpInterface(LinalgOp linalgOp)
Checks whether linalgOp conforms to ContractionOpInterface.
void populateLinalgGenericOpsSpecializationPatterns(RewritePatternSet &patterns)
Populates patterns with patterns to convert linalg.generic ops to named ops where possible.
std::optional< Value > isaFillOpInterface(GenericOp genericOp)
Checks whether genericOp is semantically equivalent to a linalg.fill.
bool isaElemwiseSingleBinaryOpInterface(GenericOp genericOp)
Checks whether genericOp is semantically equivalent to a single linalg elementwise binary op e....
detail::InFlightRemark failed(Location loc, RemarkOpts opts)
Report an optimization remark that failed.
Definition Remarks.h:573
Include the generated interface declarations.
LogicalResult applyPatternsGreedily(Region &region, const FrozenRewritePatternSet &patterns, GreedyRewriteConfig config=GreedyRewriteConfig(), bool *changed=nullptr)
Rewrite ops in the given region, which must be isolated from above, by repeatedly applying the highes...
@ DimId
Dimensional identifier.
Definition AffineExpr.h:59
llvm::SetVector< T, Vector, Set, N > SetVector
Definition LLVM.h:123
const FrozenRewritePatternSet & patterns
void getForwardSlice(Operation *op, SetVector< Operation * > *forwardSlice, const ForwardSliceOptions &options={})
Fills forwardSlice with the computed forward slice (i.e.