1 //===- LinalgTransformOps.cpp - Implementation of Linalg transform 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/Linalg/TransformOps/LinalgTransformOps.h"
10 
11 #include "mlir/Dialect/Affine/IR/AffineOps.h"
12 #include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
13 #include "mlir/Dialect/Linalg/IR/Linalg.h"
14 #include "mlir/Dialect/Linalg/Transforms/Transforms.h"
15 #include "mlir/Dialect/PDL/IR/PDL.h"
16 #include "mlir/Dialect/PDL/IR/PDLTypes.h"
17 #include "mlir/Dialect/Transform/IR/TransformDialect.h"
18 #include "mlir/Dialect/Transform/IR/TransformInterfaces.h"
19 #include "mlir/Parser/Parser.h"
20 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
21 
22 using namespace mlir;
23 using namespace mlir::linalg;
24 using namespace mlir::transform;
25 
26 /// Extracts a vector of int64_t from an array attribute. Asserts if the
27 /// attribute contains values other than integers.
28 static SmallVector<int64_t> extractI64Array(ArrayAttr attr) {
29   SmallVector<int64_t> result;
30   result.reserve(attr.size());
31   for (APInt value : attr.getAsValueRange<IntegerAttr>())
32     result.push_back(value.getSExtValue());
33   return result;
34 }
35 
36 /// Extracts a vector of unsigned from an array attribute. Asserts if the
37 /// attribute contains values other than intergers. May truncate.
38 static SmallVector<unsigned> extractUIntArray(ArrayAttr attr) {
39   SmallVector<unsigned> result;
40   result.reserve(attr.size());
41   for (APInt value : attr.getAsValueRange<IntegerAttr>())
42     result.push_back(value.getZExtValue());
43   return result;
44 }
45 
46 namespace {
47 /// A simple pattern rewriter that implements no special logic.
48 class SimpleRewriter : public PatternRewriter {
49 public:
50   SimpleRewriter(MLIRContext *context) : PatternRewriter(context) {}
51 };
52 } // namespace
53 
54 /// Attempts to apply the pattern specified as template argument to the given
55 /// operation. The pattern is expected to have a `returningMatchAndRewrite`
56 /// function that returns the "main" result or failure. Returns failure if the
57 /// pattern failed to apply. Extra arguments are forwarded to the pattern
58 /// constructor.
59 template <typename PatternTy, typename... Args>
60 static FailureOr<LinalgOp> tryApply(Operation *operation, Args &&...args) {
61   // Check if the given operation has the type expected by the pattern.
62   using OpTy = typename llvm::function_traits<
63       decltype(&PatternTy::returningMatchAndRewrite)>::template arg_t<0>;
64   auto op = dyn_cast<OpTy>(operation);
65   if (!op)
66     return failure();
67 
68   // Apply the pattern directly to the op.
69   PatternTy pattern(operation->getContext(), std::forward<Args>(args)...);
70   SimpleRewriter rewriter(operation->getContext());
71   rewriter.setInsertionPoint(operation);
72   auto result = pattern.returningMatchAndRewrite(op, rewriter);
73   if (failed(result))
74     return failure();
75   return cast<LinalgOp>(result->getOperation());
76 }
77 
78 //===----------------------------------------------------------------------===//
79 // DecomposeOp
80 //===----------------------------------------------------------------------===//
81 
82 DiagnosedSilenceableFailure
83 transform::DecomposeOp::applyToOne(linalg::LinalgOp target,
84                                    SmallVectorImpl<Operation *> &results,
85                                    transform::TransformState &state) {
86   FailureOr<LinalgOp> windowed =
87       tryApply<DownscaleSizeOneWindowed2DConvolution>(target);
88   if (succeeded(windowed)) {
89     results.push_back(*windowed);
90     return DiagnosedSilenceableFailure(success());
91   }
92   FailureOr<LinalgOp> depthwise =
93       tryApply<DownscaleDepthwiseConv2DNhwcHwcOp>(target);
94   if (succeeded(depthwise)) {
95     results.push_back(*depthwise);
96     return DiagnosedSilenceableFailure(success());
97   }
98   results.assign(1, nullptr);
99   return emitDefaultSilenceableFailure(target);
100 }
101 
102 //===----------------------------------------------------------------------===//
103 // FuseOp
104 //===----------------------------------------------------------------------===//
105 
106 /// Apply a tiling transformation to all payload ops and store both the
107 /// tiled operation as well as the created tile loops.
108 static LogicalResult
109 applyTilingToAll(Operation *transformOp, ArrayRef<Operation *> payloadOps,
110                  unsigned numLoops,
111                  transform::TransformResults &transformResults,
112                  function_ref<FailureOr<TiledLinalgOp>(LinalgOp)> applyFn) {
113   SmallVector<Operation *> tiledLinalgOps;
114   SmallVector<SmallVector<Operation *>> loopOps(numLoops);
115   for (unsigned int i = 0; i < numLoops; ++i)
116     loopOps[i].reserve(payloadOps.size());
117 
118   for (Operation *target : payloadOps) {
119     auto linalgOp = dyn_cast<linalg::LinalgOp>(target);
120     if (!linalgOp)
121       return transformOp->emitError("only LinalgOps are supported");
122 
123     FailureOr<TiledLinalgOp> tiled = applyFn(linalgOp);
124     if (failed(tiled))
125       return failure();
126 
127     tiledLinalgOps.push_back(tiled->op);
128     if (tiled->loops.size() != numLoops)
129       // Not enough loops were generated. This usually means that the input size
130       // was smaller than the tiling size.
131       // TODO: LinalgTilingPattern should return failure().
132       return failure();
133     for (unsigned int i = 0; i < numLoops; ++i)
134       loopOps[i].push_back(tiled->loops[i]);
135   }
136 
137   transformResults.set(transformOp->getOpResult(0), tiledLinalgOps);
138   for (unsigned int i = 0; i < numLoops; ++i)
139     transformResults.set(transformOp->getOpResult(i + 1), loopOps[i]);
140   return success();
141 }
142 
143 /// Parse a tiling-like operation that returns the tiled op as well as the
144 /// created tile loops. The function counts the non-zero tile sizes to compute
145 /// the number of results.
146 static ParseResult parseTileLikeOp(OpAsmParser &parser, OperationState &result,
147                                    StringRef sizesAttrName) {
148   OpAsmParser::UnresolvedOperand targetOperand;
149   SMLoc opLoc = parser.getCurrentLocation();
150   if (parser.parseOperand(targetOperand) ||
151       parser.parseOptionalAttrDict(result.attributes))
152     return failure();
153   Attribute sizesAttr = result.attributes.get(sizesAttrName);
154   if (!sizesAttr)
155     return parser.emitError(opLoc)
156            << "expected '" << sizesAttrName << "' attribute";
157   auto sizesArrayAttr = sizesAttr.dyn_cast<ArrayAttr>();
158   if (!sizesArrayAttr)
159     return parser.emitError(opLoc)
160            << "'" << sizesAttrName << "' attribute must be an array";
161   Type pdlOpType = parser.getBuilder().getType<pdl::OperationType>();
162   size_t numExpectedLoops =
163       sizesArrayAttr.size() - llvm::count(extractI64Array(sizesArrayAttr), 0);
164   result.addTypes(SmallVector<Type>(numExpectedLoops + 1, pdlOpType));
165   if (parser.resolveOperand(targetOperand, pdlOpType, result.operands))
166     return failure();
167   return success();
168 }
169 
170 DiagnosedSilenceableFailure
171 transform::FuseOp::apply(mlir::transform::TransformResults &transformResults,
172                          mlir::transform::TransformState &state) {
173   LinalgTilingAndFusionOptions fusionOptions;
174   fusionOptions.tileSizes = extractI64Array(getTileSizes());
175   fusionOptions.tileInterchange = extractI64Array(getTileInterchange());
176 
177   LogicalResult result = applyTilingToAll(
178       getOperation(), state.getPayloadOps(getTarget()),
179       fusionOptions.tileSizes.size() - llvm::count(fusionOptions.tileSizes, 0),
180       transformResults, [&](LinalgOp linalgOp) -> FailureOr<TiledLinalgOp> {
181         LinalgTileAndFuseTensorOpsPattern pattern(getContext(), fusionOptions);
182         SimpleRewriter rewriter(getContext());
183         rewriter.setInsertionPoint(linalgOp);
184         FailureOr<TileLoopNest> tileLoopNest =
185             pattern.returningMatchAndRewrite(linalgOp, rewriter);
186         if (failed(tileLoopNest))
187           return failure();
188 
189         TiledLinalgOp tiledLinalgOp;
190         tiledLinalgOp.op = tileLoopNest->getRootOp();
191         tiledLinalgOp.loops = {tileLoopNest->getLoopOps().begin(),
192                                tileLoopNest->getLoopOps().end()};
193         return tiledLinalgOp;
194       });
195   return DiagnosedSilenceableFailure(result);
196 }
197 
198 ParseResult transform::FuseOp::parse(OpAsmParser &parser,
199                                      OperationState &result) {
200   return parseTileLikeOp(
201       parser, result,
202       transform::FuseOp::getTileSizesAttrName(result.name).getValue());
203 }
204 
205 void transform::FuseOp::print(OpAsmPrinter &p) {
206   p << ' ';
207   p << getTarget();
208   p.printOptionalAttrDict((*this)->getAttrs());
209 }
210 
211 LogicalResult transform::FuseOp::verify() {
212   SmallVector<int64_t> permutation = extractI64Array(getTileInterchange());
213   auto sequence = llvm::to_vector(llvm::seq<int64_t>(0, permutation.size()));
214   if (!std::is_permutation(sequence.begin(), sequence.end(),
215                            permutation.begin(), permutation.end())) {
216     return emitOpError() << "expects interchange to be a permutation, found "
217                          << getTileInterchange();
218   }
219   return success();
220 }
221 
222 //===----------------------------------------------------------------------===//
223 // GeneralizeOp
224 //===----------------------------------------------------------------------===//
225 
226 DiagnosedSilenceableFailure
227 transform::GeneralizeOp::applyToOne(linalg::LinalgOp target,
228                                     SmallVectorImpl<Operation *> &results,
229                                     transform::TransformState &state) {
230   // Exit early if no transformation is needed.
231   if (isa<GenericOp>(target)) {
232     results.push_back(target);
233     return DiagnosedSilenceableFailure(success());
234   }
235   FailureOr<LinalgOp> generic = tryApply<LinalgGeneralizationPattern>(target);
236   if (succeeded(generic)) {
237     results.push_back(generic->getOperation());
238     return DiagnosedSilenceableFailure(success());
239   }
240   results.assign(1, nullptr);
241   return emitDefaultSilenceableFailure(target);
242 }
243 
244 //===----------------------------------------------------------------------===//
245 // InterchangeOp
246 //===----------------------------------------------------------------------===//
247 
248 DiagnosedSilenceableFailure
249 transform::InterchangeOp::applyToOne(linalg::GenericOp target,
250                                      SmallVectorImpl<Operation *> &results,
251                                      transform::TransformState &state) {
252   SmallVector<unsigned> interchangeVector =
253       extractUIntArray(getIteratorInterchange());
254   // Exit early if no transformation is needed.
255   if (interchangeVector.empty()) {
256     results.push_back(target);
257     return DiagnosedSilenceableFailure(success());
258   }
259   SimpleRewriter rewriter(target->getContext());
260   FailureOr<GenericOp> res =
261       interchangeGenericOp(rewriter, target, interchangeVector);
262   if (failed(res))
263     return DiagnosedSilenceableFailure::definiteFailure();
264   results.push_back(res->getOperation());
265   return DiagnosedSilenceableFailure(success());
266 }
267 
268 LogicalResult transform::InterchangeOp::verify() {
269   SmallVector<unsigned> permutation =
270       extractUIntArray(getIteratorInterchange());
271   auto sequence = llvm::to_vector(llvm::seq<unsigned>(0, permutation.size()));
272   if (!std::is_permutation(sequence.begin(), sequence.end(),
273                            permutation.begin(), permutation.end())) {
274     return emitOpError()
275            << "expects iterator_interchange to be a permutation, found "
276            << getIteratorInterchange();
277   }
278   return success();
279 }
280 
281 //===---------------------------------------------------------------------===//
282 // MultiTileSizesOp
283 //===---------------------------------------------------------------------===//
284 
285 DiagnosedSilenceableFailure transform::MultiTileSizesOp::applyToOne(
286     LinalgOp target, SmallVector<Operation *> &results, TransformState &state) {
287   OpBuilder builder(target.getContext());
288   builder.setInsertionPoint(target);
289   OpFoldResult targetSize = builder.getIndexAttr(getTargetSize());
290   OpFoldResult divisor = builder.getIndexAttr(getDivisor());
291   FailureOr<MultiSizeSpecification> spec = computeMultiTileSizes(
292       builder, target, getDimension(), targetSize, divisor);
293   if (failed(spec)) {
294     return emitSilenceableError() << "could not generate tile size computation";
295   }
296 
297   AffineExpr s0 = builder.getAffineSymbolExpr(0);
298   AffineExpr s1 = builder.getAffineSymbolExpr(1);
299   Operation *splitPoint =
300       makeComposedAffineApply(builder, target.getLoc(), s0 * s1,
301                               {spec->lowTileSize, spec->lowTripCount});
302   Operation *lowTileSize = spec->lowTileSize.getDefiningOp();
303   Operation *highTileSize = spec->highTileSize.getDefiningOp();
304   assert(lowTileSize && highTileSize && splitPoint &&
305          "tile sizes are not produced by operations");
306   results.reserve(results.size() + 3);
307   results.push_back(lowTileSize);
308   results.push_back(highTileSize);
309   results.push_back(splitPoint);
310   return DiagnosedSilenceableFailure::success();
311 }
312 
313 void transform::MultiTileSizesOp::getEffects(
314     SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
315   onlyReadsHandle(getTarget(), effects);
316   producesHandle(getResults(), effects);
317   modifiesPayload(effects);
318 }
319 
320 //===---------------------------------------------------------------------===//
321 // PadOp
322 //===---------------------------------------------------------------------===//
323 
324 DiagnosedSilenceableFailure
325 transform::PadOp::applyToOne(linalg::LinalgOp target,
326                              SmallVectorImpl<Operation *> &results,
327                              transform::TransformState &state) {
328   // Convert the integer packing flags to booleans.
329   SmallVector<bool> packPaddings;
330   for (int64_t packPadding : extractI64Array(getPackPaddings()))
331     packPaddings.push_back(static_cast<bool>(packPadding));
332 
333   // Convert the padding values to attributes.
334   SmallVector<Attribute> paddingValues;
335   for (auto const &it :
336        llvm::zip(getPaddingValues(), target->getOperandTypes())) {
337     Attribute attr = std::get<0>(it);
338     Type elementType = getElementTypeOrSelf(std::get<1>(it));
339     // Try to parse string attributes to obtain an attribute of element type.
340     if (auto stringAttr = attr.dyn_cast<StringAttr>()) {
341       paddingValues.push_back(
342           parseAttribute(attr.cast<StringAttr>(), elementType));
343       if (!paddingValues.back()) {
344         auto diag = this->emitOpError("expects a padding that parses to ")
345                     << elementType << ", got " << std::get<0>(it);
346         diag.attachNote(target.getLoc()) << "when applied to this op";
347         return DiagnosedSilenceableFailure::definiteFailure();
348       }
349       continue;
350     }
351     // Otherwise, add the attribute directly.
352     if (attr.getType() != elementType) {
353       auto diag = this->emitOpError("expects a padding value of type ")
354                   << elementType << ", got " << attr;
355       diag.attachNote(target.getLoc()) << "when applied to this op";
356       return DiagnosedSilenceableFailure::definiteFailure();
357     }
358     paddingValues.push_back(attr);
359   }
360 
361   // Extract the transpose vectors.
362   SmallVector<SmallVector<int64_t>> transposePaddings;
363   for (Attribute transposeVector : getTransposePaddings().cast<ArrayAttr>())
364     transposePaddings.push_back(
365         extractI64Array(transposeVector.cast<ArrayAttr>()));
366 
367   LinalgPaddingOptions paddingOptions;
368   paddingOptions.setPaddingValues(paddingValues);
369   paddingOptions.setPaddingDimensions(extractI64Array(getPaddingDimensions()));
370   paddingOptions.setPackPaddings(packPaddings);
371   paddingOptions.setHoistPaddings(extractI64Array(getHoistPaddings()));
372   paddingOptions.setTransposePaddings(transposePaddings);
373 
374   FailureOr<LinalgOp> result =
375       tryApply<LinalgPaddingPattern>(target, paddingOptions);
376   if (succeeded(result)) {
377     results.push_back(result->getOperation());
378     return DiagnosedSilenceableFailure(success());
379   }
380 
381   results.assign(1, nullptr);
382   return emitDefaultSilenceableFailure(target);
383 }
384 
385 LogicalResult transform::PadOp::verify() {
386   SmallVector<int64_t> packPaddings = extractI64Array(getPackPaddings());
387   if (any_of(packPaddings, [](int64_t packPadding) {
388         return packPadding != 0 && packPadding != 1;
389       })) {
390     return emitOpError()
391            << "expects pack_paddings to contain booleans (0/1), found "
392            << getPackPaddings();
393   }
394 
395   SmallVector<int64_t> paddingDimensions =
396       extractI64Array(getPaddingDimensions());
397   if (any_of(paddingDimensions,
398              [](int64_t paddingDimension) { return paddingDimension < 0; })) {
399     return emitOpError()
400            << "expects padding_dimensions to contain positive integers, found "
401            << getPaddingDimensions();
402   }
403 
404   SmallVector<int64_t> hoistPaddings = extractI64Array(getHoistPaddings());
405   if (any_of(hoistPaddings,
406              [](int64_t hoistPadding) { return hoistPadding < 0; })) {
407     return emitOpError()
408            << "expects hoist_paddings to contain positive integers, found "
409            << getHoistPaddings();
410   }
411 
412   ArrayAttr transposes = getTransposePaddings();
413   for (Attribute attr : transposes) {
414     SmallVector<int64_t> transpose = extractFromI64ArrayAttr(attr);
415     auto sequence = llvm::to_vector(llvm::seq<int64_t>(0, transpose.size()));
416     if (!std::is_permutation(sequence.begin(), sequence.end(),
417                              transpose.begin(), transpose.end())) {
418       return emitOpError()
419              << "expects transpose_paddings to be a permutation, found "
420              << attr;
421     }
422   }
423   return success();
424 }
425 
426 //===----------------------------------------------------------------------===//
427 // ScalarizeOp
428 //===----------------------------------------------------------------------===//
429 
430 DiagnosedSilenceableFailure
431 transform::ScalarizeOp::applyToOne(linalg::LinalgOp target,
432                                    SmallVectorImpl<Operation *> &results,
433                                    transform::TransformState &state) {
434   LinalgTilingOptions tilingOptions;
435   tilingOptions.scalarizeDynamicDims();
436   // Tiling with "scalarize_dyn_dims" actually sets the same lambda as the tile
437   // sizes and asserts that it is not already set.
438   SmallVector<int64_t> emptyTileSizes;
439   LinalgTilingPattern pattern(getContext(), tilingOptions);
440   SimpleRewriter rewriter(getContext());
441   rewriter.setInsertionPoint(target);
442   FailureOr<TiledLinalgOp> result =
443       pattern.returningMatchAndRewrite(target, rewriter);
444   if (failed(result))
445     return DiagnosedSilenceableFailure(reportUnknownTransformError(target));
446 
447   results.push_back(result->op);
448   return DiagnosedSilenceableFailure(success());
449 }
450 
451 //===----------------------------------------------------------------------===//
452 // SplitOp
453 //===----------------------------------------------------------------------===//
454 
455 DiagnosedSilenceableFailure SplitOp::apply(TransformResults &results,
456                                            TransformState &state) {
457   // Collect the dynamic split points if provided.
458   ArrayRef<Operation *> payload = state.getPayloadOps(getTarget());
459   SimpleRewriter rewriter(getContext());
460   SmallVector<OpFoldResult> splitPoints;
461   splitPoints.reserve(payload.size());
462   if (getDynamicSplitPoint()) {
463     auto diag = DiagnosedSilenceableFailure::success();
464     splitPoints = llvm::to_vector(llvm::map_range(
465         state.getPayloadOps(getDynamicSplitPoint()), [&](Operation *op) {
466           if (op->getNumResults() != 1 ||
467               !op->getResult(0).getType().isIndex()) {
468             diag = emitSilenceableError()
469                    << "expected dynamic split point handle to point to a "
470                       "single-result index-typed op";
471             diag.attachNote(op->getLoc()) << "dynamic split point";
472           }
473           return OpFoldResult(op->getResult(0));
474         }));
475     if (!diag.succeeded())
476       return diag;
477 
478     if (splitPoints.size() != payload.size()) {
479       emitError() << "expected the dynamic split point handle to point to as "
480                      "many operations ("
481                   << splitPoints.size() << ") as the target handle ("
482                   << payload.size() << ")";
483       return DiagnosedSilenceableFailure::definiteFailure();
484     }
485   } else {
486     splitPoints.resize(payload.size(),
487                        rewriter.getIndexAttr(getStaticSplitPoint()));
488   }
489 
490   // Split each target operation.
491   SmallVector<Operation *> first, second;
492   for (const auto &pair : llvm::zip(payload, splitPoints)) {
493     Operation *target = std::get<0>(pair);
494     auto linalgOp = dyn_cast<LinalgOp>(target);
495     if (!linalgOp) {
496       auto diag = emitSilenceableError() << "only applies to structured ops";
497       diag.attachNote(target->getLoc()) << "target op";
498       return diag;
499     }
500 
501     if (getDimension() >= linalgOp.getNumLoops()) {
502       auto diag = emitSilenceableError() << "dimension " << getDimension()
503                                          << " does not exist in target op";
504       diag.attachNote(target->getLoc()) << "target op";
505       return diag;
506     }
507 
508     rewriter.setInsertionPoint(linalgOp);
509     std::tie(first.emplace_back(), second.emplace_back()) =
510         linalg::splitOp(rewriter, linalgOp, getDimension(), std::get<1>(pair));
511   }
512 
513   results.set(getFirst().cast<OpResult>(), first);
514   results.set(getSecond().cast<OpResult>(), second);
515   return DiagnosedSilenceableFailure::success();
516 }
517 
518 void SplitOp::getEffects(
519     SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
520   consumesHandle(getTarget(), effects);
521   if (getDynamicSplitPoint())
522     onlyReadsHandle(getDynamicSplitPoint(), effects);
523   producesHandle(getResults(), effects);
524   modifiesPayload(effects);
525 }
526 
527 ParseResult SplitOp::parse(OpAsmParser &parser, OperationState &result) {
528   OpAsmParser::UnresolvedOperand target, dynamicSplitPoint;
529   IntegerAttr staticSplitPoint;
530   auto pdlOperationType =
531       pdl::OperationType::get(parser.getBuilder().getContext());
532   if (parser.parseOperand(target) ||
533       parser.resolveOperand(target, pdlOperationType, result.operands) ||
534       parser.parseKeyword("after"))
535     return failure();
536 
537   OptionalParseResult dynamicPointParseResult =
538       parser.parseOptionalOperand(dynamicSplitPoint);
539   if (!dynamicPointParseResult.hasValue()) {
540     int64_t staticSplitPointValue;
541     if (failed(parser.parseInteger(staticSplitPointValue)))
542       return failure();
543 
544     staticSplitPoint =
545         parser.getBuilder().getI64IntegerAttr(staticSplitPointValue);
546   } else {
547     if (failed(*dynamicPointParseResult) ||
548         parser.resolveOperand(dynamicSplitPoint, pdlOperationType,
549                               result.operands)) {
550       return failure();
551     }
552 
553     staticSplitPoint =
554         parser.getBuilder().getI64IntegerAttr(ShapedType::kDynamicSize);
555   }
556 
557   result.addAttribute(
558       SplitOp::getStaticSplitPointAttrName(result.name).getValue(),
559       staticSplitPoint);
560   if (failed(parser.parseOptionalAttrDict(result.attributes)))
561     return failure();
562 
563   result.addTypes({pdlOperationType, pdlOperationType});
564   return success();
565 }
566 
567 void SplitOp::print(OpAsmPrinter &printer) {
568   printer << " " << getTarget() << " after ";
569   int64_t staticSplitSize = static_cast<int64_t>(getStaticSplitPoint());
570   if (staticSplitSize != ShapedType::kDynamicSize)
571     printer << staticSplitSize;
572   else
573     printer << getDynamicSplitPoint();
574   printer << " ";
575   printer.printOptionalAttrDict(getOperation()->getAttrs(),
576                                 {getStaticSplitPointAttrName()});
577 }
578 
579 LogicalResult SplitOp::verify() {
580   if ((static_cast<int64_t>(getStaticSplitPoint()) !=
581        ShapedType::kDynamicSize) ^
582       (getDynamicSplitPoint() == nullptr)) {
583     return emitOpError()
584            << "expects either a dynamic or a static split point to be provided";
585   }
586   return success();
587 }
588 
589 //===----------------------------------------------------------------------===//
590 // SplitReductionOp
591 //===----------------------------------------------------------------------===//
592 
593 DiagnosedSilenceableFailure
594 transform::SplitReductionOp::applyToOne(linalg::LinalgOp target,
595                                         SmallVectorImpl<Operation *> &results,
596                                         transform::TransformState &state) {
597   ControlSplitReductionFn splitFn = [&](LinalgOp) {
598     return std::pair<int64_t, unsigned>(getSplitFactor(),
599                                         getInsertSplitDimension());
600   };
601   SimpleRewriter rewriter(getContext());
602   rewriter.setInsertionPoint(target);
603   FailureOr<SplitReductionResult> splitResult =
604       (getUseScalingAlgorithm())
605           ? splitReductionByScaling(rewriter, target, splitFn, getUseAlloc())
606           : splitReduction(rewriter, target, splitFn, getUseAlloc());
607   if (failed(splitResult))
608     return DiagnosedSilenceableFailure(reportUnknownTransformError(target));
609 
610   results.push_back(splitResult->initOrAlloc);
611   results.push_back(splitResult->fillOp);
612   results.push_back(splitResult->splitLinalgOp);
613   results.push_back(splitResult->resultCombiningLinalgOp);
614   return DiagnosedSilenceableFailure(success());
615 }
616 
617 //===----------------------------------------------------------------------===//
618 // TileOp
619 //===----------------------------------------------------------------------===//
620 
621 DiagnosedSilenceableFailure
622 transform::TileOp::apply(TransformResults &transformResults,
623                          TransformState &state) {
624   LinalgTilingOptions tilingOptions;
625   SmallVector<int64_t> tileSizes = extractI64Array(getStaticSizes());
626 
627   ArrayRef<Operation *> targets = state.getPayloadOps(getTarget());
628   SmallVector<ArrayRef<Operation *>> dynamicSizeProducers;
629   dynamicSizeProducers.reserve(getDynamicSizes().size());
630   for (Value dynamicSizeProducerHandle : getDynamicSizes()) {
631     dynamicSizeProducers.push_back(
632         state.getPayloadOps(dynamicSizeProducerHandle));
633 
634     if (dynamicSizeProducers.back().size() != targets.size()) {
635       DiagnosedSilenceableFailure diag =
636           emitSilenceableError()
637           << "expected as many dynamic size-producing operations ("
638           << dynamicSizeProducers.back().size() << ") as target ops ("
639           << targets.size() << ")";
640       diag.attachNote(dynamicSizeProducerHandle.getLoc()) << "for this handle";
641       return diag;
642     }
643 
644     for (Operation *op : dynamicSizeProducers.back()) {
645       if (op->getNumResults() == 1 &&
646           op->getResult(0).getType().isa<IndexType>())
647         continue;
648       DiagnosedSilenceableFailure diag =
649           emitSilenceableError() << "expected sizes to be produced by ops "
650                                     "with a single index-type result";
651       diag.attachNote(op->getLoc()) << "size producer op";
652       diag.attachNote(dynamicSizeProducerHandle.getLoc()) << "for this handle";
653       return diag;
654     }
655   }
656 
657   SmallVector<Operation *> tiled;
658   SmallVector<SmallVector<Operation *, 4>, 4> loops;
659   loops.resize(getLoops().size());
660   for (auto &en : llvm::enumerate(targets)) {
661     auto linalgOp = dyn_cast<LinalgOp>(en.value());
662     if (!linalgOp) {
663       DiagnosedSilenceableFailure diag = emitSilenceableError()
664                                          << "only linalg ops are supported";
665       diag.attachNote(en.value()->getLoc()) << "target op";
666       return diag;
667     }
668 
669     unsigned index = en.index();
670     if (!tileSizes.empty()) {
671       tilingOptions.setTileSizeComputationFunction(
672           [&, index](OpBuilder &b, Operation *) {
673             SmallVector<Value, 4> sizes;
674             sizes.reserve(tileSizes.size());
675             unsigned dynamicIdx = 0;
676             for (OpFoldResult ofr : getMixedSizes()) {
677               if (auto attr = ofr.dyn_cast<Attribute>()) {
678                 sizes.push_back(b.create<arith::ConstantIndexOp>(
679                     getLoc(), attr.cast<IntegerAttr>().getInt()));
680               } else {
681                 sizes.push_back(
682                     dynamicSizeProducers[dynamicIdx++][index]->getResult(0));
683               }
684             }
685             return sizes;
686           });
687     }
688 
689     tilingOptions.setInterchange(extractUIntArray(getInterchange()));
690     LinalgTilingPattern pattern(getContext(), tilingOptions);
691     SimpleRewriter rewriter(linalgOp.getContext());
692     FailureOr<TiledLinalgOp> tiledOp =
693         pattern.returningMatchAndRewrite(linalgOp, rewriter);
694     if (failed(tiledOp))
695       return DiagnosedSilenceableFailure::definiteFailure();
696 
697     tiled.push_back(tiledOp->op);
698     for (const auto &en2 : llvm::enumerate(tiledOp->loops))
699       loops[en2.index()].push_back(en2.value());
700   }
701 
702   transformResults.set(getTiledLinalgOp().cast<OpResult>(), tiled);
703   for (const auto &en : llvm::enumerate(loops))
704     transformResults.set(getLoops()[en.index()].cast<OpResult>(), en.value());
705 
706   return DiagnosedSilenceableFailure::success();
707 }
708 
709 SmallVector<OpFoldResult> transform::TileOp::getMixedSizes() {
710   ValueRange dynamic = getDynamicSizes();
711   SmallVector<int64_t> tileSizes = extractI64Array(getStaticSizes());
712   SmallVector<OpFoldResult> results;
713   results.reserve(tileSizes.size());
714   unsigned dynamicPos = 0;
715   Builder builder(getContext());
716   for (int64_t size : tileSizes) {
717     if (size == ShapedType::kDynamicSize) {
718       results.push_back(dynamic[dynamicPos++]);
719     } else {
720       results.push_back(builder.getIndexAttr(size));
721     }
722   }
723   return results;
724 }
725 
726 ParseResult transform::TileOp::parse(OpAsmParser &parser,
727                                      OperationState &result) {
728   OpAsmParser::UnresolvedOperand target;
729   SmallVector<OpAsmParser::UnresolvedOperand> dynamicSizes;
730   ArrayAttr staticSizes;
731   auto pdlOperationType = pdl::OperationType::get(parser.getContext());
732   if (parser.parseOperand(target) ||
733       parser.resolveOperand(target, pdlOperationType, result.operands) ||
734       parseOperandsOrIntegersSizesList(parser, dynamicSizes, staticSizes) ||
735       parser.resolveOperands(dynamicSizes, pdlOperationType, result.operands) ||
736       parser.parseOptionalAttrDict(result.attributes))
737     return ParseResult::failure();
738 
739   result.addAttribute(getStaticSizesAttrName(result.name), staticSizes);
740   size_t numExpectedLoops =
741       staticSizes.size() - llvm::count(extractI64Array(staticSizes), 0);
742   result.addTypes(SmallVector<Type>(numExpectedLoops + 1, pdlOperationType));
743   return success();
744 }
745 
746 void TileOp::print(OpAsmPrinter &p) {
747   p << ' ' << getTarget();
748   printOperandsOrIntegersSizesList(p, getOperation(), getDynamicSizes(),
749                                    getStaticSizes());
750   p.printOptionalAttrDict((*this)->getAttrs(), {getStaticSizesAttrName()});
751 }
752 
753 void transform::TileOp::getEffects(
754     SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
755   consumesHandle(getTarget(), effects);
756   onlyReadsHandle(getDynamicSizes(), effects);
757   producesHandle(getTiledLinalgOp(), effects);
758   producesHandle(getLoops(), effects);
759   modifiesPayload(effects);
760 }
761 
762 //===----------------------------------------------------------------------===//
763 // VectorizeOp
764 //===----------------------------------------------------------------------===//
765 
766 DiagnosedSilenceableFailure
767 transform::VectorizeOp::applyToOne(Operation *target,
768                                    SmallVectorImpl<Operation *> &results,
769                                    transform::TransformState &state) {
770   if (!target->hasTrait<OpTrait::IsIsolatedFromAbove>()) {
771     auto diag = this->emitOpError("requires isolated-from-above targets");
772     diag.attachNote(target->getLoc()) << "non-isolated target";
773     return DiagnosedSilenceableFailure::definiteFailure();
774   }
775 
776   MLIRContext *ctx = getContext();
777   RewritePatternSet patterns(ctx);
778   patterns.add<LinalgVectorizationPattern>(ctx);
779 
780   vector::populateVectorTransferPermutationMapLoweringPatterns(patterns);
781   vector::populateVectorReductionToContractPatterns(patterns);
782   patterns.add<linalg::LinalgCopyVTRForwardingPattern,
783                linalg::LinalgCopyVTWForwardingPattern>(ctx,
784                                                        /*benefit=*/2);
785   vector::TransferReadOp::getCanonicalizationPatterns(patterns, ctx);
786   vector::TransferWriteOp::getCanonicalizationPatterns(patterns, ctx);
787   if (getVectorizePadding())
788     linalg::populatePadOpVectorizationPatterns(patterns);
789 
790   if (failed(applyPatternsAndFoldGreedily(target, std::move(patterns))))
791     return DiagnosedSilenceableFailure(reportUnknownTransformError(target));
792 
793   results.push_back(target);
794   return DiagnosedSilenceableFailure(success());
795 }
796 
797 //===----------------------------------------------------------------------===//
798 // Transform op registration
799 //===----------------------------------------------------------------------===//
800 
801 namespace {
802 /// Registers new ops and declares PDL as dependent dialect since the additional
803 /// ops are using PDL types for operands and results.
804 class LinalgTransformDialectExtension
805     : public transform::TransformDialectExtension<
806           LinalgTransformDialectExtension> {
807 public:
808   LinalgTransformDialectExtension() {
809     declareDependentDialect<AffineDialect>();
810     declareDependentDialect<arith::ArithmeticDialect>();
811     declareDependentDialect<pdl::PDLDialect>();
812     declareDependentDialect<scf::SCFDialect>();
813     declareDependentDialect<vector::VectorDialect>();
814     registerTransformOps<
815 #define GET_OP_LIST
816 #include "mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp.inc"
817         >();
818   }
819 };
820 } // namespace
821 
822 #define GET_OP_CLASSES
823 #include "mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp.inc"
824 
825 void mlir::linalg::registerTransformDialectExtension(
826     DialectRegistry &registry) {
827   registry.addExtensions<LinalgTransformDialectExtension>();
828 }
829