MLIR  22.0.0git
Transforms.h
Go to the documentation of this file.
1 //===- Transforms.h - Linalg transformations as patterns --------*- C++ -*-===//
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 #ifndef MLIR_DIALECT_LINALG_TRANSFORMS_TRANSFORMS_H
10 #define MLIR_DIALECT_LINALG_TRANSFORMS_TRANSFORMS_H
11 
12 #include <utility>
13 
23 #include "mlir/IR/OpDefinition.h"
24 #include "mlir/IR/PatternMatch.h"
27 #include "llvm/ADT/SmallBitVector.h"
28 #include "llvm/ADT/SmallSet.h"
29 
30 namespace mlir {
31 namespace bufferization {
32 class AllocTensorOp;
33 class OneShotAnalysisState;
34 class BufferizationState;
35 } // namespace bufferization
36 
37 namespace linalg {
38 
39 class LinalgOp;
40 enum class WinogradConv2DFmr : uint32_t;
41 
42 //===----------------------------------------------------------------------===//
43 // Utils.
44 //===----------------------------------------------------------------------===//
45 
46 /// Return vector::CombiningKind for the given op.
47 std::optional<vector::CombiningKind> getCombinerOpKind(Operation *combinerOp);
48 
49 //===----------------------------------------------------------------------===//
50 // Bufferization-related transforms.
51 //===----------------------------------------------------------------------===//
52 
54  enum class AllocOp { MemrefAlloc = 0, MemrefAlloca = 1 };
55  AllocOp allocOp = AllocOp::MemrefAlloc;
56 
57  enum class MemcpyOp {
58  MaterializeInDestination = 0,
59  MemrefCopy = 1,
60  LinalgCopy = 2
61  };
62  MemcpyOp memcpyOp = MemcpyOp::MaterializeInDestination;
63 
64  /// If set to "true", only the destination tensor operands are bufferized to
65  /// a new allocation (and wrapped in "bufferization.to_tensor"), but not the
66  /// targeted op itself.
67  bool bufferizeDestinationOnly = false;
68 
69  /// If set to "true", a `memref.dealloc` operation will be emitted for each
70  /// allocated buffer. Otherwise, the memory is leaked, which is useful if
71  /// the buffer deallocation pipeline should be run after bufferization is
72  /// done.
73  bool emitDealloc = false;
74 };
75 
76 /// Materialize a buffer allocation for the given tensor.pad op and lower the
77 /// op to linalg.fill/linalg.generic + bufferization.materialize_in_destination.
78 /// E.g.:
79 ///
80 /// %0 = tensor.pad low[%l] high[%h] %t ...
81 ///
82 /// is lowered to:
83 ///
84 /// %alloc = memref.alloc
85 /// linalg.fill ... outs(%alloc)
86 /// %subview = memref.subview %alloc [%l] [...] [1]
87 /// bufferization.materialize_in_destination %t in %subview
88 /// %0 = bufferization.to_tensor %alloc restrict writable
89 ///
90 /// In addition to rewriting the IR as shown above, this function returns the
91 /// newly allocated buffer. The `insertionPoint` parameter can be used to
92 /// specify a custom insertion point for the buffer allocation.
95  tensor::PadOp padOp, Attribute memorySpace = {},
96  Operation *insertionPoint = nullptr);
97 
98 /// Materialize a buffer allocation for the given vector.mask op and bufferize
99 /// the op, including its region. E.g.:
100 ///
101 /// %0 = vector.mask {
102 /// vector.transfer_write %v, %t : vector<16xf32>, tensor<?xf32>
103 /// } : vector<16xi1> -> tensor<?xf32>
104 ///
105 /// is lowered to:
106 ///
107 /// %alloc = memref.alloc
108 /// bufferization.materialize_in_destination %t in %subview
109 /// vector.mask {
110 /// vector.transfer_write %arg0, %alloc : vector<16xf32>, memref<?xf32>
111 /// } : vector<16xi1>
112 /// %0 = bufferization.to_tensor %alloc restrict writable
113 ///
114 /// In addition to rewriting the IR as shown above, this function returns the
115 /// newly allocated buffer. The `insertionPoint` parameter can be used to
116 /// specify a custom insertion point for the buffer allocation.
118  const BufferizeToAllocationOptions &options,
119  vector::MaskOp maskOp, Attribute memorySpace = {},
120  Operation *insertionPoint = nullptr);
121 
122 /// Materialize a buffer allocation for the given bufferization.alloc_tensor op
123 /// and lower the op to memref.alloc + memref.tensor_store.
124 ///
125 /// In addition to rewriting the IR, this function returns the newly allocated
126 /// buffer. The `insertionPoint` parameter can be used to specify a custom
127 /// insertion point for the buffer allocation.
128 Value bufferizeToAllocation(RewriterBase &rewriter,
129  const BufferizeToAllocationOptions &options,
130  bufferization::AllocTensorOp allocTensorOp,
131  Attribute memorySpace = {},
132  Operation *insertionPoint = nullptr);
133 
134 /// Bufferize the given op with tensor semantics and materialize the result in
135 /// a newly allocated buffer.
136 ///
137 /// Only bufferizable ops that bufferize to a memory write or have an
138 /// aliasing OpOperand (and do not themselves bufferize to an allocation) are
139 /// supported. They are bufferized using their BufferizableOpInterface
140 /// implementation.
141 ///
142 /// Selected ops that bufferize to an allocation (or need special handling) are
143 /// also supported:
144 /// - tensor.pad
145 /// - vector.mask
146 ///
147 /// This function returns the newly allocated buffer. The `insertionPoint`
148 /// parameter can be used to specify a custom insertion point for the buffer
149 /// allocation.
150 Value bufferizeToAllocation(RewriterBase &rewriter,
151  const BufferizeToAllocationOptions &options,
152  Operation *op, Attribute memorySpace = {},
153  Operation *insertionPoint = nullptr);
154 
155 /// Try to eliminate tensor::EmptyOps inside `op` that are anchored on a
156 /// LinalgOp. This transforms looks for LinalgOps that have an unused output
157 /// operand and an input operand that is rooted in a tensor::EmptyOp. The
158 /// tensor::EmptyOp uses are replaced with the output operand and the two
159 /// operands of the LinalgOp are swapped.
160 ///
161 /// Example:
162 /// %0 = tensor.empty()
163 /// %1 = linalg.matmul ins(...) outs(%0)
164 /// %2 = linalg.generic ins(%1) outs(%dest) {
165 /// ^bb0(%in: f32, %out: f32):
166 /// // out not used
167 /// }
168 ///
169 /// The IR is transformed as follows:
170 /// %0 = tensor.empty()
171 /// %1 = linalg.matmul ins(...) outs(%dest)
172 /// %2 = linalg.generic ins(%0) outs(%1) {
173 /// ^bb0(%in: f32, %out: f32):
174 /// // Use %out instead of %in
175 /// }
176 ///
177 /// The "ins" operand has no uses inside the body of the LinalgOp and can be
178 /// folded away with existing cleanup patterns. Afterwards, the tensor::EmptyOp
179 /// can also fold away.
181  RewriterBase &rewriter, Operation *op,
182  bufferization::OneShotAnalysisState &state);
183 
184 //===----------------------------------------------------------------------===//
185 // Structs that configure the behavior of various transformations.
186 //===----------------------------------------------------------------------===//
187 
189  std::function<SmallVector<Value, 4>(OpBuilder &, Operation *)>;
190 
192  /// Computation function that returns the tile sizes for each operation.
193  /// Delayed construction of constant tile sizes should occur to interoperate
194  /// with folding.
196 
199  tileSizeComputationFunction = std::move(fun);
200  return *this;
201  }
202  /// Set the `tileSizeComputationFunction` to return the values `ts`. The
203  /// values must not fold away when tiling. Otherwise, use a more robust
204  /// `tileSizeComputationFunction`.
206  tileSizeComputationFunction = [=](OpBuilder &, Operation *) { return ts; };
207  return *this;
208  }
209  /// Convenience function to set the `tileSizeComputationFunction` to a
210  /// function that computes tile sizes at the point they are needed. Allows
211  /// proper interaction with folding.
213 
214  /// Tile all dynamic dimensions by 1. I.e., scalarize those dimensions.
215  /// Note: `scalarizeDynamicDims` and `setTileSizes` cannot be used together.
217 
218  /// The interchange vector to reorder the tiled loops.
220 
222  interchangeVector.assign(interchange.begin(), interchange.end());
223  return *this;
224  }
225 
226  /// The type of tile loops to generate.
228 
230  loopType = lt;
231  return *this;
232  }
233 
234  /// When specified, specifies distribution of generated tile loops to
235  /// processors.
236  std::optional<LinalgLoopDistributionOptions> distribution;
237 
240  distribution = std::move(distributionOptions);
241  return *this;
242  }
243 
244  /// Specification markers of how to distribute the `linalg.tiled_loop`.
246 
248  distributionTypes.assign(types.begin(), types.end());
249  return *this;
250  }
251 
252  /// Peel the specified loops.
254 
256  peeledLoops.clear();
257  peeledLoops.append(loops.begin(), loops.end());
258  return *this;
259  }
260 };
261 
263  /// Tile sizes used to tile the root operation.
266  tileSizes.assign(ts.begin(), ts.end());
267  return *this;
268  }
269  /// Tile interchange used to permute the tile loops.
271  /// When specified, specifies distribution of generated tile loops to
272  /// processors.
273  std::optional<LinalgLoopDistributionOptions> tileDistribution;
276  tileDistribution = std::move(distributionOptions);
277  return *this;
278  }
279 };
280 
282  /// A padding value for every operand.
285  paddingValues.assign(pv.begin(), pv.end());
286  return *this;
287  }
288  /// A list of iterator dimensions to pad.
291  paddingDimensions.assign(pd.begin(), pd.end());
292  return *this;
293  }
294  /// A list of multiples to which each padding dimension should be padded to.
295  std::optional<SmallVector<int64_t>> padToMultipleOf;
297  padToMultipleOf.emplace(m.begin(), m.end());
298  return *this;
299  }
300  /// A mapping between an operand and shape dim, and a size for a padding
301  /// dimension. Each size is expected to be greater or equal than the
302  /// corresponding shape dim. If no value is provided then the constant upper
303  /// bound will be used.
305  LinalgPaddingOptions &setSizeToPadTo(unsigned operandIndex, unsigned dimIndex,
306  OpFoldResult size) {
307  assert(size && "expected non-null size");
308  sizeToPadTo[{operandIndex, dimIndex}] = size;
309  return *this;
310  }
311  /// Given the operand index and shape dim it returns the size to pad to.
312  OpFoldResult getSizeToPadTo(unsigned operandIndex, unsigned dimIndex) const {
313  return sizeToPadTo.lookup_or(
314  std::pair<unsigned, unsigned>(operandIndex, dimIndex), nullptr);
315  }
316 
317  /// A flag for every operand to mark the PadOp as nofold which enables
318  /// packing for statically shaped operands.
321  nofoldFlags.assign(pp.begin(), pp.end());
322  return *this;
323  }
324  /// A number of loops to hoist the PadOp out for every operand.
327  hoistPaddings.assign(hp.begin(), hp.end());
328  return *this;
329  }
330  /// A permutation vector for every operand used to transpose the packed
331  /// PadOp results.
335  transposePaddings.assign(tp.begin(), tp.end());
336  return *this;
337  }
338  enum class CopyBackOp : int8_t {
339  None = 0,
341  LinalgCopy = 2
342  };
343  /// The op to be used for copying the padded result to the original
344  /// destination tensor.
347  copyBackOp = op;
348  return *this;
349  }
350 };
351 
353  /// A padding value for every operand.
356  paddingValues.assign(pv.begin(), pv.end());
357  return *this;
358  }
359  /// A list of iterator dimensions sizes to pad to.
362  paddingSizes.assign(m.begin(), m.end());
363  return *this;
364  }
365  /// Pad iterator `paddingDimension[i]` to next multiple of `paddingSizes[i]`
366  /// if true. Otherwise pad to `paddingSizes[i]`.
369  padToMultipleOf = b;
370  return *this;
371  }
372 };
373 
374 /// Callback function type used to perform the allocation for the promoted
375 /// `subView`. In `boundingSubViewsize` a best attempt is made to find the
376 /// smallest constant value for the size of the buffer needed for each
377 /// dimension. If that is not possible, contains the dynamic size of the
378 /// subview. The call back should return the buffer to use.
379 using AllocBufferCallbackFn = std::function<std::optional<Value>(
380  OpBuilder &b, memref::SubViewOp subView,
381  ArrayRef<Value> boundingSubViewSize, DataLayout &layout)>;
382 
383 /// Callback function type used to deallocate the buffers used to hold the
384 /// promoted subview.
386  std::function<LogicalResult(OpBuilder &b, Value buffer)>;
387 
388 /// Callback function type used to insert copy from original subview to
389 /// subview of the promoted region for the read operands/subview of promoted
390 /// region to original subview for the results. The copy has to happen from
391 /// `src` to `dst`.
393  std::function<LogicalResult(OpBuilder &b, Value src, Value dst)>;
394 
396  /// Indices of subViews to promote. If `std::nullopt`, try to promote all
397  /// operands.
398  std::optional<DenseSet<unsigned>> operandsToPromote;
401  operandsToPromote->insert_range(operands);
402  return *this;
403  }
404  /// If ith element of `useFullTiles` is true the full view should be used
405  /// for the promoted buffer of the ith operand in `operandsToPromote`.
406  /// Otherwise the partial view will be used. The decision is defaulted to
407  /// `useFullTileBuffersDefault` when `useFullTileBuffers` is std::nullopt and
408  /// for operands missing from `useFullTileBuffers`.
409  std::optional<llvm::SmallBitVector> useFullTileBuffers;
411  unsigned size = useFullTiles.size();
412  llvm::SmallBitVector tmp(size, false);
413  for (unsigned i = 0; i < size; ++i)
414  tmp[i] = useFullTiles[i];
415  useFullTileBuffers = tmp;
416  return *this;
417  }
418  /// If true all operands unspecified by `useFullTileBuffers` will use the
419  /// full view, otherwise the partial view.
423  return *this;
424  }
425  /// If true, buffers will be allocated with the original subview size. This
426  /// may result in more dynamic allocations, in case of dynamic sizes.
429  useOriginalSubviewSize = originalSize;
430  return *this;
431  }
432  /// Alignment of promoted buffer. If `std::nullopt` do not specify alignment.
433  std::optional<unsigned> alignment;
435  alignment = align;
436  return *this;
437  }
438  /// Memory space of promoted buffer. If `std::nullopt` do not specify memory
439  /// space.
440  std::optional<Attribute> memorySpace;
442  memorySpace = memorySpc;
443  return *this;
444  }
445  /// Use alloca with the default allocation scheme.
446  bool useAlloca = false;
448  useAlloca = use;
449  return *this;
450  }
451  /// Callback function to do the allocation of the promoted buffer. If
452  /// std::nullopt, then the default allocation scheme of allocating a
453  /// memref<?xi8> buffer followed by a view operation is used.
454  std::optional<AllocBufferCallbackFn> allocationFn;
455  std::optional<DeallocBufferCallbackFn> deallocationFn;
458  DeallocBufferCallbackFn const &deallocFn) {
459  allocationFn = allocFn;
460  deallocationFn = deallocFn;
461  return *this;
462  }
463  /// Callback function to do the copy of data to and from the promoted
464  /// subview. If std::nullopt then a memref.copy is used.
465  std::optional<CopyCallbackFn> copyInFn;
466  std::optional<CopyCallbackFn> copyOutFn;
468  CopyCallbackFn const &copyOut) {
469  copyInFn = copyIn;
470  copyOutFn = copyOut;
471  return *this;
472  }
473 };
474 
475 /// Split Reduction options.
477  // Ratio used to split the reduction dimension. If the ratio is <= 1,
478  // nothing will be done.
479  int64_t ratio = 0;
480  // Index where the extra dimension is added to the intermediate tensor
481  // shape.
482  unsigned index = 0;
483  // If the inner dimension after splitting is parallel or reduction.
484  bool innerParallel = false;
485 };
486 
487 /// Function signature to control reduction splitting. This returns
488 /// `SplitReductionOptions`.
489 // TODO: don't use unsigned unless doing bit manipulation.
491  std::function<SplitReductionOptions(LinalgOp op)>;
492 
493 //===----------------------------------------------------------------------===//
494 // Preconditions that ensure the corresponding transformation succeeds and can
495 // be applied as a rewrite pattern.
496 //===----------------------------------------------------------------------===//
497 
498 /// Return true if two `linalg.generic` operations with producer/consumer
499 /// relationship through `fusedOperand` can be fused using elementwise op
500 /// fusion.
501 bool areElementwiseOpsFusable(OpOperand *fusedOperand);
502 
503 /// Promote memref.subviews feeding linalg-on-buffers operations.
504 LogicalResult promoteSubviewsPrecondition(Operation *op,
506 
507 /// Return success if the operation can be vectorized.
508 LogicalResult vectorizeOpPrecondition(Operation *op,
509  ArrayRef<int64_t> inputVectorSizes = {},
510  ArrayRef<bool> inputScalableVecDims = {},
511  bool vectorizeNDExtract = false,
512  bool flatten1DDepthwiseConv = false);
513 
514 //===----------------------------------------------------------------------===//
515 // Transformations exposed as functional-style API calls.
516 //===----------------------------------------------------------------------===//
517 
519 
520 /// Transformation to drop unit-extent dimensions from `linalg.generic`
521 /// operations.
524 
527 
528  using ControlFnTy = std::function<SmallVector<unsigned>(Operation *)>;
530  if (auto genericOp = dyn_cast_or_null<GenericOp>(op)) {
531  return llvm::to_vector(llvm::seq<unsigned>(0, genericOp.getNumLoops()));
532  }
533  if (auto padOp = dyn_cast_or_null<tensor::PadOp>(op)) {
534  return llvm::to_vector(
535  llvm::seq<unsigned>(0, padOp.getSourceType().getRank()));
536  }
537  return SmallVector<unsigned>{};
538  };
539 };
541  linalg::GenericOp resultOp;
543 };
544 FailureOr<DropUnitDimsResult> dropUnitDims(RewriterBase &rewriter,
545  GenericOp genericOp,
547 
548 /// Fuse two `linalg.generic` operations that have a producer-consumer
549 /// relationship captured through `fusedOperand`. The method expects
550 /// that `areElementwiseOpsFusable` returns true for the given `fusedOperand`.
554 };
555 FailureOr<ElementwiseOpFusionResult>
556 fuseElementwiseOps(RewriterBase &rewriter, OpOperand *fusedOperand);
557 
558 /// Returns a set of indices of the producer's results which would
559 /// be preserved after the fusion.
560 /// * There is a chance that the implementation of the transformation does not
561 /// agree with the result of this method. This function gives a prediction based
562 /// on an optimized fusion.
563 llvm::SmallDenseSet<int> getPreservedProducerResults(GenericOp producer,
564  GenericOp consumer,
565  OpOperand *fusedOperand);
566 
567 /// Try to peel and canonicalize loop `op` and return the new result.
568 /// Also applies affine_min/max bounds simplification on the fly where relevant.
569 // TODO: Add support for scf.parallel and affine.for loops.
571 
572 /// Peel 'loops' and applies affine_min/max bounds simplification on the fly
573 /// where relevant.
574 void peelLoops(RewriterBase &rewriter, ArrayRef<scf::ForOp> loops);
575 
576 /// Pad the iterator dimensions `options.paddingDimensions` of all `opToPad`
577 /// operands to a static bounding box. The original `opToPad` is cloned and
578 /// operates on the padded tensors.
579 ///
580 /// * "options.padToMultipleOf" indicates that each padding dimension should be
581 /// padded to the specified multiple.
582 /// * Use "options.paddingValues" and "options.nofoldFlags" to set padding
583 /// value and nofold attribute of the created tensor::PadOps, respectively.
584 /// * The unpadded results (extracted slice of the cloned operation) are
585 /// returned via `replacements`.
586 /// * The tensor::PadOps are returned via `padOps`.
587 /// * "options.copyBackOp" specifies the op type for copying back the unpadded
588 /// result to the original destination tensor.
589 LogicalResult rewriteAsPaddedOp(RewriterBase &rewriter, LinalgOp opToPad,
591  LinalgOp &paddedOp,
592  SmallVector<Value> &replacements,
594 
595 /// Helper function to compute the padded shape of the given value `v` of
596 /// `RankedTensorType` given:
597 /// - the `indexingSizes` as a list of OpFoldResult.
598 /// - an `indexingMap` that encodes how the padded shape varies with
599 /// increases in `indexingSizes`.
600 /// The implementation iteratively combines increases from contributing using
601 /// affine.apply operations.
602 /// The `indexingMap` + `indexingSizes` encoding suits StructuredOps and
603 /// provides a gentle portability path for Linalg-like ops with affine maps.
604 /// In the future, more general interfaces can be devised to encode similar
605 /// shape evolutions and map between an op and its operands.
608  AffineMap indexingMap, ArrayRef<OpFoldResult> indexingSizes,
610 
612  std::function<FailureOr<SmallVector<OpFoldResult>>(
614  const PadTilingInterfaceOptions &)>;
615 
616 /// Specific helper for Linalg ops.
617 FailureOr<SmallVector<OpFoldResult>> computeIndexingMapOpInterfacePaddedShape(
618  RewriterBase &rewriter, OpOperand &operandToPad,
619  ArrayRef<Range> iterationDomain, const PadTilingInterfaceOptions &options);
620 
621 /// Pad the iterator dimensions `options.paddingDimensions` of `opToPad`.
622 ///
623 /// * "options.paddingSizes" indicates that each padding dimension should be
624 /// padded to the specified padding size.
625 /// * "options.padToMultipleOf" indicates that the paddingSizes should be
626 // interpreted as the bounding box (dynamic) value to pad to.
627 /// * Use "options.paddingValues" to set the padding value of the created
628 // tensor::PadOp.
629 /// * The tensor::PadOp is returned on success.
630 
631 FailureOr<TilingInterface>
632 rewriteAsPaddedOp(RewriterBase &rewriter, TilingInterface opToPad,
633  const PadTilingInterfaceOptions &constOptions,
635  PadSizeComputationFunction computePaddingSizeFun =
637 
638 namespace detail {
639 
640 /// Helper struct to hold the results of building a packing loop nest.
644  TransposeOp maybeTransposeOp;
645  tensor::PadOp hoistedPadOp;
646 };
647 
648 /// Build the packing loop nest required to hoist `opToHoist` above
649 /// `outermostEnclosingForOp`.
650 /// The loop nest is built just before `outermostEnclosingForOp`.
651 FailureOr<PackingResult>
652 buildPackingLoopNest(RewriterBase &rewriter, tensor::PadOp opToHoist,
653  scf::ForOp outermostEnclosingForOp,
654  ArrayRef<int64_t> transposeVector);
655 
656 } // namespace detail
657 
658 /// Mechanically hoist padding operations on tensors by `numLoops` into a new,
659 /// generally larger tensor. This achieves packing of multiple padding ops into
660 /// a larger tensor. On success, `opToHoist` is replaced by the cloned version
661 /// in the packing loop so the caller can continue reasoning about the padding
662 /// operation. If `transposeVector` is non-empty, hoist padding introduces a
663 /// TransposeOp to transpose the padded tensor before inserting it into the
664 /// packed tensor. A `transposeVector` can change the storage order of the
665 /// padded tensor but does not change the order of the pack or compute loops.
666 ///
667 /// TODO: In the future, we should consider rewriting as a linalg.pack after
668 /// hoisting since this abstraction is now available.
669 ///
670 /// Example in pseudo-mlir:
671 /// =======================
672 ///
673 /// If hoistPaddingOnTensors is called with `nLoops` = 2 on the following IR.
674 /// ```
675 /// scf.for (%i, %j, %k)
676 /// %st0 = tensor.extract_slice f(%i, %k) : ... to tensor<?x?xf32>
677 /// %0 = tensor.pad %st0 low[0, 0] high[...] {
678 /// ^bb0( ... ):
679 /// linalg.yield %pad
680 /// } : tensor<?x?xf32> to tensor<4x8xf32>
681 /// compute(%0)
682 /// ```
683 ///
684 /// IR resembling the following is produced:
685 ///
686 /// ```
687 /// scf.for (%i) {
688 /// %packed_init = tensor.empty range(%j) : tensor<?x4x8xf32>
689 /// %packed = scf.for (%k) iter_args(%p : %packed_init) {
690 /// %st0 = tensor.extract_slice f(%i, %k) : ... to tensor<?x?xf32>
691 /// %0 = tensor.pad %st0 low[0, 0] high[...] {
692 /// ^bb0( ... ):
693 /// linalg.yield %pad
694 /// } : tensor<?x?xf32> to tensor<4x8xf32>
695 /// %1 = tensor.insert_slice %0 ...
696 /// : tensor<4x8xf32> to tensor<?x4x8xf32>
697 /// scf.yield %1: tensor<?x4x8xf32>
698 /// } -> tensor<?x4x8xf32>
699 /// scf.for (%j, %k) {
700 /// %st0 = tensor.extract_slice %packed [%k, 0, 0][1, 4, 8][1, 1, 1] :
701 /// tensor<?x4x8xf32> to tensor<4x8xf32>
702 /// compute(%st0)
703 /// }
704 /// }
705 /// ```
706 FailureOr<Value>
707 hoistPaddingOnTensors(RewriterBase &rewriter, tensor::PadOp opToHoist,
708  int64_t numLoops, ArrayRef<int64_t> transposeVector,
709  tensor::PadOp &hoistedOp,
710  SmallVectorImpl<TransposeOp> &transposeOps);
711 /// Calls into `hoistPaddingOnTensors` with a local IRRewriter.
712 FailureOr<Value>
713 hoistPaddingOnTensors(tensor::PadOp opToHoist, int64_t numLoops,
714  ArrayRef<int64_t> transposeVector,
715  tensor::PadOp &hoistedOp,
716  SmallVectorImpl<TransposeOp> &transposeOps);
717 
718 /// Apply padding and hoisting to `linalgOp` according to the configuration
719 /// specified in `options`.
720 FailureOr<LinalgOp> padAndHoistLinalgOp(RewriterBase &rewriter,
721  LinalgOp linalgOp,
723 
724 /// Split the given `op` into two parts along the given iteration space
725 /// `dimension` at the specified `splitPoint`, and return the two parts.
726 /// If the second part is statically known to be empty, do not create it
727 /// and return nullptr instead. Error state is signalled by returning
728 /// a pair of nullptrs.
729 ///
730 /// For example, the following op:
731 ///
732 /// linalg.matmul ins(%0, %1 : tensor<128x32xf32>, tensor<32x64xf32>)
733 /// outs(%2 : tensor<128x64xf32>)
734 ///
735 /// split along the first dimension at position 42 will result in:
736 ///
737 /// %3 = tensor.extract_slice %0[0, 0][42, 32][1, 1]
738 /// %4 = tensor.extract_slice %2[0, 0][42, 64][1, 1]
739 /// %5 = linalg.matmul ins(%3, %1 : tensor<42x32xf32>, tensor<32x64xf32>)
740 /// outs(%5 : tensor<42x64xf32>)
741 /// %6 = tensor.insert_slice %5 into %2[0, 0][42, 64][1, 1]
742 ///
743 /// %7 = tensor.extract_slice %0[42, 0][86, 32][1, 1]
744 /// %8 = tensor.extract_slice %6[42, 0][86, 64][1, 1]
745 /// %9 = linalg.matmul ins(%7, %1 : tensor<86x32xf32>, tensor<32x64xf32>)
746 /// outs(%8 : tensor<86x64xf32>)
747 /// tensor.insert_slice %5 into %6[42, 0][86, 64][1, 1]
748 ///
749 /// Note that there is no simplification other than constant propagation applied
750 /// to slice extraction and insertion.
751 std::pair<TilingInterface, TilingInterface> splitOp(RewriterBase &rewriter,
752  TilingInterface op,
753  unsigned dimension,
754  OpFoldResult splitPoint);
755 
756 /// Perform standalone tiling of a single LinalgOp by `tileSizes`.
757 /// and permute the loop nest according to `interchangeVector`
758 /// The permutation is expressed as a list of integers that specify
759 /// the new ordering of the loop nest. The length of `interchangeVector`
760 /// must be equal to the length of `tileSizes`.
761 /// An empty vector is interpreted as the identity permutation and the
762 /// transformation returns early.
763 ///
764 /// Return a struct containing the tiled loops in the specified order
765 /// and the cloned op if successful, std::nullopt otherwise.
766 ///
767 /// E.g. the permutation `(i,j,k) -> (j,k,i)` is expressed by
768 /// `interchangeVector = [1,2,0]`. All values in `interchangeVector` must be
769 /// integers, in the range 0..`tileSizes.size()` without duplications
770 /// (i.e. `[1,1,2]` is an invalid permutation).
772  LinalgOp op;
775 };
776 FailureOr<TiledLinalgOp> tileLinalgOp(RewriterBase &b, LinalgOp op,
778 
779 /// Interchange the `iterator_types` and `iterator_maps` dimensions and adapts
780 /// the index accesses of `op`. This is an in-place transformation controlled
781 /// by `interchangeVector`. An empty vector is interpreted as the identity
782 /// permutation and the transformation returns early.
783 ///
784 /// E.g. the permutation `(i,j,k) -> (j,k,i)` is expressed with
785 /// `interchangeVector = [1,2,0]`. All values in `interchangeVector` must be
786 /// integers, in the range 0..`op.rank` without duplications
787 /// (i.e. `[1,1,2]` is an invalid permutation).
788 ///
789 /// Return failure if the permutation is not valid.
790 FailureOr<GenericOp> interchangeGenericOp(RewriterBase &rewriter,
791  GenericOp genericOp,
792  ArrayRef<unsigned> interchangeVector);
793 
794 /// Create a GenericOp from the given named operation `linalgOp` and replace
795 /// the given `linalgOp`.
796 /// Return failure if `linalgOp` is a GenericOp or misses a region builder.
797 FailureOr<GenericOp> generalizeNamedOp(RewriterBase &rewriter,
798  LinalgOp linalgOp);
799 
800 /// Create a namedOp from the given GenericOp and replace the GenericOp.
801 /// Currently we can specialize only trivial linalg copy operations.
802 FailureOr<LinalgOp> specializeGenericOp(RewriterBase &rewriter,
803  GenericOp genericOp);
804 
805 /// Create a new buffer using the `allocationFn` provided. The size of this
806 /// buffer is either the original subview size when 'useOriginalSubviewSize' is
807 /// set to true or the smallest constant bounding size along each dimension that
808 /// can be computed for the size of the result of `subView`. Returns the
809 /// allocated buffer as `fullLocalView` and the view that matches the size of
810 /// the result of subview operation as `partialLocalView`.
814 };
815 FailureOr<PromotionInfo>
816 promoteSubviewAsNewBuffer(OpBuilder &b, Location loc, memref::SubViewOp subView,
817  bool useOriginalSubviewSize,
818  const AllocBufferCallbackFn &allocationFn,
819  DataLayout &layout);
820 
821 /// Promote the `subViews` into a new buffer allocated at the insertion point
822 /// `b`. Promotion occurs in 3 steps:
823 /// 1. Create a new buffer for a full tile (i.e. not clipped at the
824 /// boundary).
825 /// 2. Take a full view on the buffer.
826 /// 3. Take a partial slice of the full view in step 2. and copy into it.
827 ///
828 /// Return the modified linalg op (the modification happens in place) as well
829 /// as all the copy ops created.
830 FailureOr<LinalgOp> promoteSubViews(OpBuilder &b, LinalgOp op,
832 
833 /// Allocate the subview in the GPU workgroup memory.
834 std::optional<Value> allocateWorkgroupMemory(OpBuilder &builder,
835  memref::SubViewOp subview,
836  ArrayRef<Value> sizeBounds,
837  DataLayout &);
838 
839 /// In case of GPU group memory there is no need to deallocate.
840 LogicalResult deallocateWorkgroupMemory(OpBuilder &, Value /*buffer*/);
841 
842 /// Create Memref copy operations and add gpu barrier guards before and after
843 /// the copy operation to ensure data integrity.
844 LogicalResult copyToWorkgroupMemory(OpBuilder &b, Value src, Value dst);
845 
846 /// Allocate the subview in the GPU private memory.
847 std::optional<Value> allocateGPUPrivateMemory(OpBuilder &builder,
848  memref::SubViewOp subview,
849  ArrayRef<Value> sizeBounds,
850  DataLayout &);
851 
852 /// Normal copy to between src and dst.
853 LogicalResult copyToGPUPrivateMemory(OpBuilder &b, Value src, Value dst);
854 
855 /// In case of GPU private memory there is no need to deallocate since the
856 /// memory is freed when going outside of the scope.
857 LogicalResult deallocateGPUPrivateMemory(OpBuilder &, Value /*buffer*/);
858 
859 /// Return true if there's dedicated logic in the Linalg Vectorizer to
860 /// vectorize this Op, false otherwise.
861 ///
862 /// Note that this helper merely implements a very high level check and that the
863 /// vectorizer also requires various additional pre-conditions to be met for it
864 /// to work (these are checked by the vectorizer itself).
866 
867 /// Transformation information returned after vectorizing.
869  /// Results of the vectorization transform to replace the original operation.
871 };
872 /// Returns a `VectorizationResult` containing the results of the vectorized op,
873 /// or failure if the transformation fails. If provided, `inputVectorSizes` are
874 /// used to vectorize this operation. `inputVectorSizes` must match the rank of
875 /// the iteration space of the operation and the input vector sizes must be
876 /// greater than or equal to their counterpart iteration space sizes, if static.
877 /// `inputVectorShapes` also allows the vectorization of operations with dynamic
878 /// shapes.
879 FailureOr<VectorizationResult>
880 vectorize(RewriterBase &rewriter, Operation *op,
881  ArrayRef<int64_t> inputVectorSizes = {},
882  ArrayRef<bool> inputScalableVecDims = {},
883  bool vectorizeNDExtract = false, bool flatten1DDepthwiseConv = false,
884  bool assumeDynamicDimsMatchVecSizes = false);
885 
886 /// Emit a suitable vector form for a Copy op with fully static shape.
887 LogicalResult vectorizeCopy(RewriterBase &builder, memref::CopyOp copyOp);
888 
889 /// Emit a loop nest of `scf.for` with the proper body for `linalgOp`.
890 FailureOr<LinalgLoops> linalgOpToLoops(RewriterBase &rewriter,
891  LinalgOp linalgOp);
892 
893 /// Emit a loop nest of `scf.parallel` with the proper body for `linalgOp`.
894 FailureOr<LinalgLoops> linalgOpToParallelLoops(RewriterBase &rewriter,
895  LinalgOp linalgOp);
896 
897 /// Emit a loop nest of `affine.for` with the proper body for `linalgOp`.
898 FailureOr<LinalgLoops> linalgOpToAffineLoops(RewriterBase &rewriter,
899  LinalgOp linalgOp);
900 
901 /// Creates a number of ranges equal to the number of non-zero in `tileSizes`.
902 /// One for each loop of the LinalgOp that is tiled. The `tileSizes` argument
903 /// has one entry per surrounding loop. It uses zero as the convention that a
904 /// particular loop is not tiled. This convention simplifies implementations
905 /// by avoiding affine map manipulations. The returned ranges correspond to
906 /// the loop ranges, in the proper order, that are tiled and for which new
907 /// loops will be created. Also the function returns a map from loop indices
908 /// of the LinalgOp to the corresponding non-empty range indices of newly
909 /// created loops.
911 std::tuple<SmallVector<Range, 4>, LoopIndexToRangeIndexMap>
913  ArrayRef<OpFoldResult> allShapeSizes,
914  ArrayRef<OpFoldResult> allTileSizes);
915 
916 namespace detail {
917 template <typename T>
919  /// Tile sizes.
921  /// Number of tiles associated with each size.
923 };
924 
925 template <typename T>
927  /// Tile sizes.
929  /// Number of tiles associated with each size.
931 };
932 
933 } // namespace detail
934 
935 /// A description of a multi-size tiling comprising tile sizes and numbers of
936 /// tiles, expressed as Values which may or may not be constant. Multi-size
937 /// currently means two-size.
939  : public detail::MultiSizeSpecificationBase<Value> {};
941  : public detail::MultiSizeSpecificationBase<int64_t> {};
942 
946  : public detail::ContinuousTileSizeSpecificationBase<int64_t> {};
947 
948 /// Emits the IR computing the multi-sized tiling specification with two tile
949 /// sizes not exceeding `targetSize`, each divisible by `sizeDivisor`, such
950 /// that there exist numbers of tiles with these sizes that fully cover the
951 /// given iteration space `dimension` of the structured `op`.
952 ///
953 /// The computation is as follows:
954 ///
955 /// b = originalTripCount floordiv sizeDivisor
956 /// t = (targetSize + sizeDivisor - 1) floordiv sizeDivisor
957 /// d = (b + t - 1) floordiv t
958 /// s = (b floordiv d) * sizeDivisor
959 /// v = b % d
960 /// u = d - v
961 ///
962 /// where the tile sizes are `s` and `s` + `sizeDivisor`, and the numbers of
963 /// the corresponding tiles are `u` and `v`, respectively. Alternatively,
964 ///
965 /// s * u + (s + sizeDivisor) * v == original size,
966 /// where s mod sizeDivisor = 0.
967 ///
968 /// Expects all values to be positive. In some cases with the target tile size
969 /// sufficiently close to the dimension shape and non-unit divisor, it is
970 /// impossible to compute such sizes. If `emitAssertion` is set, also emit the
971 /// assertion that size computation succeeded.
972 ///
973 /// Returns the specification consisting of both tile values and the number of
974 /// tiles of each size.
975 FailureOr<MultiSizeSpecification>
976 computeMultiTileSizes(OpBuilder &builder, LinalgOp op, unsigned dimension,
977  OpFoldResult targetSize, OpFoldResult divisor,
978  bool emitAssertions = true);
979 FailureOr<StaticMultiSizeSpecification>
980 computeStaticMultiTileSizes(LinalgOp op, unsigned dimension, int64_t targetSize,
981  int64_t divisor);
982 
983 FailureOr<StaticContinuousTileSizeSpecification>
984 computeStaticContinuousTileSizes(LinalgOp op, unsigned dimension,
985  unsigned targetSize);
986 FailureOr<ContinuousTileSizeSpecification>
987 computeContinuousTileSizes(OpBuilder &builder, TilingInterface op,
988  unsigned dimension, OpFoldResult targetSize,
989  bool emitAssertions);
990 
991 /// Transformation information returned after reduction tiling.
993  /// The partial reduction tiled op generated.
995  /// The final reduction operation merging all the partial reductions.
997  /// Initial values used for partial reductions.
999  /// The `scf.forall` operation that iterate over the tiles.
1000  scf::ForallOp loops;
1001 };
1002 
1003 /// Method to tile a reduction to parallel iterations computing partial
1004 /// reductions. After the loop all the partial reduction are merged into a final
1005 /// reduction. For example for the following sequence
1006 ///
1007 /// ```mlir
1008 /// %0 = linalg.generic %in ["parallel", "reduction"]
1009 /// : tensor<7x9xf32> -> tensor<7xf32>
1010 /// ```
1011 ///
1012 /// into:
1013 ///
1014 /// ```mlir
1015 /// %0 = linalg.fill ... : tensor<7x4xf32>
1016 /// %1 = scf.forall (%iv) in (%c4) shared_outs(%arg0 = %0)
1017 /// -> (tensor<7x4xf32>) {
1018 /// %2 = tensor.extract_slice %arg3 : tensor<7x4xf32> to tensor<7xf32>
1019 /// %3 = tensor.extract_slice %in : tensor<7x9xf32> -> tensor<7x?xf32>
1020 /// %4 = linalg.generic %2, %3 ["parallel", "reduction"]
1021 /// : tensor<7x?xf32> -> tensor<7xf32>
1022 /// %5 = tensor.insert_slice %3, %arg0[0, %iv] : tensor<7x4xf32>
1023 /// }
1024 /// %6 = linalg.generic %1 ["parallel", "reduction"]
1025 /// : tensor<7x4xf32> -> tensor<7xf32>
1026 /// ```
1027 FailureOr<ForallReductionTilingResult>
1028 tileReductionUsingForall(RewriterBase &b, PartialReductionOpInterface op,
1029  ArrayRef<OpFoldResult> numThreads,
1030  ArrayRef<OpFoldResult> tileSizes = {},
1031  std::optional<ArrayAttr> mapping = std::nullopt);
1032 
1033 /// All indices returned by IndexOp should be invariant with respect to
1034 /// tiling. Therefore, if an operation is tiled, we have to transform the
1035 /// indices accordingly, i.e. offset them by the values of the corresponding
1036 /// induction variables that are captured implicitly in the body of the op.
1037 ///
1038 /// Example. `linalg.generic` before tiling:
1039 ///
1040 /// #id_2d = (i, j) -> (i, j)
1041 /// #pointwise_2d_trait = {
1042 /// indexing_maps = [#id_2d, #id_2d],
1043 /// iterator_types = ["parallel", "parallel"]
1044 /// }
1045 /// linalg.generic #pointwise_2d_trait %operand, %result {
1046 /// ^bb0(%operand_in: f32, %result_in: f32):
1047 /// %i = linalg.index 0 : index
1048 /// %j = linalg.index 1 : index
1049 /// <some operations that use %i, %j>
1050 /// }: memref<50x100xf32>, memref<50x100xf32>
1051 ///
1052 /// After tiling pass with tiles sizes 10 and 25:
1053 ///
1054 /// #strided = (i, j)[s0, s1, s2] -> (i * s1 + s0 + j * s2)
1055 ///
1056 /// %c1 = arith.constant 1 : index
1057 /// %c0 = arith.constant 0 : index
1058 /// %c25 = arith.constant 25 : index
1059 /// %c10 = arith.constant 10 : index
1060 /// operand_dim_0 = dim %operand, 0 : memref<50x100xf32>
1061 /// operand_dim_1 = dim %operand, 1 : memref<50x100xf32>
1062 /// scf.for %k = %c0 to operand_dim_0 step %c10 {
1063 /// scf.for %l = %c0 to operand_dim_1 step %c25 {
1064 /// %4 = memref.subview %operand[%k, %l][%c10, %c25][%c1, %c1]
1065 /// : memref<50x100xf32> to memref<?x?xf32, #strided>
1066 /// %5 = memref.subview %result[%k, %l][%c10, %c25][%c1, %c1]
1067 /// : memref<50x100xf32> to memref<?x?xf32, #strided>
1068 /// linalg.generic pointwise_2d_trait %4, %5 {
1069 /// ^bb0(%operand_in: f32, %result_in: f32):
1070 /// %i = linalg.index 0 : index
1071 /// %j = linalg.index 1 : index
1072 /// // Indices `k` and `l` are implicitly captured in the body.
1073 /// %transformed_i = arith.addi %i, %k : index // index `i` is offset by
1074 /// %k %transformed_j = arith.addi %j, %l : index // index `j` is offset
1075 /// by %l
1076 /// // Every use of %i, %j is replaced with %transformed_i,
1077 /// %transformed_j <some operations that use %transformed_i,
1078 /// %transformed_j>
1079 /// }: memref<?x?xf32, #strided>, memref<?x?xf32, #strided>
1080 /// }
1081 /// }
1082 ///
1083 /// TODO: Investigate whether mixing implicit and explicit indices
1084 /// does not lead to losing information.
1085 void transformIndexOps(RewriterBase &b, LinalgOp op,
1087  const LoopIndexToRangeIndexMap &loopIndexToRangeIndex);
1088 
1089 /// Apply transformation to split the single linalg op reduction into a
1090 /// parallel and reduction dimension. Then create a new linalg.generic op
1091 /// doing the rest of the reduction. Return the new linalg op with an extra
1092 /// parallel dimension or failure if the transformation didn't happen.
1093 ///
1094 /// Example:
1095 /// ```
1096 /// %r = linalg.generic {indexing_maps = [affine_map<(d0) -> (d0)>,
1097 /// affine_map<(d0) -> ()>],
1098 /// iterator_types = ["reduction"]}
1099 /// ins(%in : tensor<32xf32>)
1100 /// outs(%out : tensor<f32>) {
1101 /// ^bb0(%arg1: f32, %arg2: f32):
1102 /// %y = arith.addf %arg1, %arg2 : f32
1103 /// linalg.yield %y : f32
1104 /// } -> tensor<f32>
1105 /// ```
1106 /// To:
1107 /// ```
1108 /// %cst = arith.constant 0.000000e+00 : f32
1109 /// %0 = tensor.expand_shape %in [[0, 1]]: tensor<32xf32> into tensor<4x8xf32>
1110 /// %1 = tensor.empty [4] : tensor<4xf32>
1111 /// %2 = linalg.fill ins(%cst : f32)
1112 /// outs(%1 : tensor<4xf32>) -> tensor<4xf32>
1113 /// %3 = linalg.generic {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>,
1114 /// affine_map<(d0, d1) -> (d0)>],
1115 /// iterator_types = ["parallel", "reduction"]}
1116 /// ins(%0 : tensor<4x8xf32>) outs(%2 : tensor<4xf32>) {
1117 /// ^bb0(%arg3: f32, %arg5: f32):
1118 /// %5 = arith.addf %arg3, %arg4 : f32
1119 /// linalg.yield %5 : f32
1120 /// } -> tensor<4xf32>
1121 /// %r = linalg.generic {indexing_maps = [affine_map<(d0) -> (d0)>,
1122 /// affine_map<(d0) -> ()>],
1123 /// iterator_types = ["reduction"]}
1124 /// ins(%3 : tensor<4xf32>) outs(%out : tensor<f32>) {
1125 /// ^bb0(%arg3: f32, %arg4: f32):
1126 /// %5 = arith.addf %arg3, %arg4 : f32
1127 /// linalg.yield %5 : f32
1128 /// } -> tensor<f32>
1129 /// ```
1132  FillOp fillOp;
1133  LinalgOp splitLinalgOp;
1135 };
1136 FailureOr<SplitReductionResult>
1137 splitReduction(RewriterBase &b, LinalgOp op,
1138  const ControlSplitReductionFn &controlSplitReductionFn,
1139  bool useAlloc = false);
1140 
1141 /// Scaling-based implementation of the split reduction transformation.
1142 /// Instead of introducing an ExpandShapeOp, this rewrites a reduction
1143 /// dimension `k` into `k * scale + kk`.
1144 ///
1145 /// Example:
1146 /// ```
1147 /// %0 = linalg.matmul ins(%A, %B: tensor<16x256xf32>, tensor<256x32xf32>)
1148 /// outs(%C: tensor<16x32xf32>) -> tensor<16x32xf32>
1149 /// ```
1150 ///
1151 /// Is transformed to:
1152 ///
1153 /// ```
1154 /// #map0 = affine_map<(d0, d1, d2, d3) -> (d0, d2 * 4 + d3)>
1155 /// #map1 = affine_map<(d0, d1, d2, d3) -> (d2 * 4 + d3, d1)>
1156 /// #map2 = affine_map<(d0, d1, d2, d3) -> (d2, d3)>
1157 /// #map3 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>
1158 /// #map4 = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
1159 /// #map5 = affine_map<(d0, d1, d2) -> (d0, d1)>
1160 /// %0 = tensor.empty [16, 32, 64] : tensor<16x32x64xf32>
1161 /// %cst = arith.constant 0.000000e+00 : f32
1162 /// %1 = linalg.fill ins(%cst : f32) outs(%0 : tensor<16x32x64xf32>) ->
1163 /// tensor<16x32x64xf32>
1164 /// %2 = tensor.empty [64, 4] : tensor<64x4xi1>
1165 ///
1166 /// %3 = linalg.generic {indexing_maps = [#map0, #map1, #map2, #map3],
1167 /// iterator_types = ["parallel", "parallel", "parallel", "reduction"]}
1168 /// ins(%A, %B, %2 : tensor<16x256xf32>, tensor<256x32xf32>,
1169 /// tensor<64x4xi1>)
1170 /// outs(%1 : tensor<16x32x64xf32>) {
1171 /// ^bb0(%arg3: f32, %arg4: f32, %arg5: i1, %arg6: f32):
1172 /// %5 = arith.mulf %arg3, %arg4 : f32
1173 /// %6 = arith.addf %arg6, %5 : f32
1174 /// linalg.yield %6 : f32
1175 /// } -> tensor<16x32x64xf32>
1176 ///
1177 /// %4 = linalg.generic {indexing_maps = [#map4, #map5],
1178 /// iterator_types = ["parallel", "parallel", "reduction"]}
1179 // ins(%3 : tensor<16x32x64xf32>)
1180 /// outs(%C : tensor<16x32xf32>) {
1181 /// ^bb0(%arg3: f32, %arg4: f32):
1182 /// %5 = arith.addf %arg3, %arg4 : f32
1183 /// linalg.yield %5 : f32
1184 /// } -> tensor<16x32xf32>
1185 ///
1186 /// return %4 : tensor<16x32xf32>
1187 /// ```
1188 FailureOr<SplitReductionResult>
1189 splitReductionByScaling(RewriterBase &b, LinalgOp op,
1190  const ControlSplitReductionFn &controlSplitReductionFn,
1191  bool useAlloc = false);
1192 
1193 /// Return `true` if a given sequence of dimensions are contiguous in the
1194 /// range of the specified indexing map.
1196 /// Return `true` if all sequences of dimensions specified in `dimSequences` are
1197 /// contiguous in all the ranges of the `maps`.
1199  ArrayRef<ReassociationIndices> dimSequences);
1200 
1203  LinalgOp collapsedOp;
1204 };
1205 
1206 /// Collapses dimensions of linalg.generic/linalg.copy operation. A precondition
1207 /// to calling this method is that for each list in `foldedIterationDim`, the
1208 /// sequence of dimensions is contiguous in domains of all `indexing_maps` of
1209 /// the `linalgOp`. This can be checked using `areDimSequencePreserved` method.
1210 /// When valid, the method also collapses the operands of the op. Returns
1211 /// replacement values of the results of the original `linalgOp` by inserting
1212 /// reshapes to get back values of compatible types.
1213 FailureOr<CollapseResult>
1214 collapseOpIterationDims(LinalgOp op,
1215  ArrayRef<ReassociationIndices> foldedIterationDims,
1216  RewriterBase &rewriter);
1217 
1219  tensor::PadOp padOp;
1220  tensor::ExpandShapeOp expandShapeOp;
1221  linalg::TransposeOp transposeOp;
1222 };
1223 
1224 /// Rewrite pack as pad + reshape + transpose.
1225 FailureOr<LowerPackResult> lowerPack(RewriterBase &rewriter,
1226  linalg::PackOp packOp,
1227  bool lowerPadLikeWithInsertSlice = true);
1228 
1230  tensor::EmptyOp emptyOp;
1231  linalg::TransposeOp transposeOp;
1232  tensor::CollapseShapeOp collapseShapeOp;
1233  tensor::ExtractSliceOp extractSliceOp;
1234 };
1235 
1236 /// Rewrite pack as empty + transpose + reshape + extract_slice.
1237 FailureOr<LowerUnPackOpResult>
1238 lowerUnPack(RewriterBase &rewriter, linalg::UnPackOp unPackOp,
1239  bool lowerUnpadLikeWithExtractSlice = true);
1240 
1241 /// Struct to hold the result of a `pack` call.
1242 struct PackResult {
1244  linalg::LinalgOp packedLinalgOp;
1246 };
1247 /// Implement packing of a single LinalgOp by `packedSizes`.
1248 /// There must be one packedSizes entry per `linalgOp` iterator.
1249 /// Return the packed Linalg op on success, failure otherwise.
1250 FailureOr<PackResult> pack(RewriterBase &rewriter, linalg::LinalgOp linalgOp,
1251  ArrayRef<OpFoldResult> packedSizes);
1252 
1253 /// Struct to hold the result of a `packTranspose` call.
1255  linalg::PackOp transposedPackOp;
1256  linalg::LinalgOp transposedLinalgOp;
1257  linalg::UnPackOp transposedUnPackOp;
1258 };
1259 /// Transpose a single PackOp -> LinalgOp -> UnPackOp chain and return the
1260 /// transposed PackOp -> LinalgOp -> UnPackOp chain after replacements.
1261 /// Return failure if either:
1262 /// 1. the `packOp` does not have the `linalgOp` as its unique use.
1263 /// 2. the `maybeUnPackOp`, if specified must be a consumer of the result tied
1264 /// to the unique `packOp` use.
1265 /// 3. `outerPerm` (resp. `innerPerm`) must be valid permutations of
1266 /// `packOp.getOuterDimsPerm` (resp. `packOp.getInnerDimsPerm`) or empty.
1267 FailureOr<PackTransposeResult>
1268 packTranspose(RewriterBase &rewriter, linalg::PackOp packOp,
1269  linalg::LinalgOp linalgOp, linalg::UnPackOp maybeUnPackOp,
1270  ArrayRef<int64_t> outerPerm, ArrayRef<int64_t> innerPerm);
1271 
1272 /// Pack a LinalgOp by greedily inferring matmul dimensions (m, n, k) where m
1273 /// and n are proper parallel dimensions and k is a proper reduction
1274 /// dimension. Packing occurs by rewriting the op as a linalg.generic and
1275 /// calling linalg::pack by `mnkPackedSizes`. The order of the packed
1276 /// dimensions is customizable: the `mnkOrder` is a permutation of {0, 1, 2}
1277 /// to reorder {m, n, k} into one of the 8 possible forms. The outer
1278 /// dimensions of the operands are not permuted at this time, this is left for
1279 /// future work.
1280 FailureOr<PackResult>
1281 packMatmulGreedily(RewriterBase &rewriter, LinalgOp linalgOp,
1282  ArrayRef<OpFoldResult> mnkPackedSizes,
1283  ArrayRef<int64_t> mnkPaddedSizesNextMultipleOf,
1284  ArrayRef<int64_t> mnkOrder);
1285 
1287  /// Minor block factors (mb, nb, kb) for packing relayout where mb, mn are
1288  /// the parallel dimensions and kb is the reduction dimension.
1290 
1291  /// If true, allows packing of dimensions that only partially fit into the
1292  /// block factors.
1293  bool allowPadding = true;
1294 
1295  /// Next multiples of the packing sizes.
1297 
1298  /// Permutation of matmul (M, N, K) dimensions order.
1300 
1301  /// Transpose LHS outer block layout [MB][KB] -> [KB][MB].
1303 
1304  /// Transpose LHS inner block layout [mb][kb] -> [kb][mb].
1306 
1307  /// Transpose RHS outer block layout [KB][NB] -> [NB][KB].
1309 
1310  /// Transpose RHS inner block layout [kb][nb] -> [nb][kb].
1312 };
1313 
1314 /// Function type which is used to control matmul packing.
1315 /// It is expected to return valid packing configuration for each operation.
1316 /// Lack of packing options indicates that no valid configuration could be
1317 /// assigned and the operation will not be packed.
1319  std::function<std::optional<BlockPackMatmulOptions>(linalg::LinalgOp)>;
1320 
1321 /// Pack a matmul operation into blocked 4D layout.
1322 ///
1323 /// Relayout a matmul operation into blocked layout with two levels of
1324 /// subdivision:
1325 /// - major 2D blocks - outer dimensions, consist of minor blocks
1326 /// - minor 2D blocks - inner dimensions, consist of scalar elements
1327 ///
1328 /// A 2D matmul MxNxK gets reshaped into blocked 4D representation
1329 /// as: [MB][NB][mb][nb] += [MB][KB][mb][kb] * [NB][KB][nb][kb]
1330 /// where the (MB, NB, KB) dimensions represent the major blocks,
1331 /// and the (mb, nb, kb) are the minor blocks of their respective
1332 /// original 2D dimensions (M, N, K).
1333 ///
1334 /// Depending on the initial operands' data layout and the specified
1335 /// packing options, the major blocks dimensions might get transposed
1336 /// e.g., [MB][KB] -> [KB][MB]. The minor blocks can also be transposed
1337 /// e.g., [mb][kb] -> [kb][mb].
1338 /// Any present batch dimensions remain unchanged.
1339 /// The final result is unpacked back to the original shape.
1340 ///
1341 /// Return failure if no valid packing options are provided.
1342 FailureOr<PackResult>
1343 blockPackMatmul(RewriterBase &rewriter, linalg::LinalgOp linalgOp,
1344  const ControlBlockPackMatmulFn &controlPackMatmul);
1345 
1346 /// Rewrite tensor.from_elements to linalg.generic.
1347 FailureOr<Operation *>
1349  tensor::FromElementsOp fromElementsOp);
1350 
1351 /// Rewrite tensor.generate to linalg.generic.
1352 FailureOr<Operation *>
1354  tensor::GenerateOp generateOp);
1355 
1356 /// Rewrite tensor.pad to linalg.generic + tensor.insert_slice.
1357 FailureOr<Operation *> rewriteInDestinationPassingStyle(RewriterBase &rewriter,
1358  tensor::PadOp padOp);
1359 
1360 /// Convert linalg.conv_2d_nhwc_hwcf into linalg.generic (for img2col packing)
1361 /// and linalg.matmul.
1362 ///
1363 /// A convolution operation can be written as a matrix-matrix multiplication by
1364 /// unfolding the cross-correlation between input and filter and explicitly copy
1365 /// overlapped sliding window inputs.
1366 ///
1367 /// Consider 2D input X with single channel input and output and 2x2 filter W:
1368 /// [x(0, 0) , x(0, 1) , ..., x(0, n) ]
1369 /// [x(1, 0) , x(1, 1) , ..., x(1, n) ]
1370 /// [. , . ,. , . ] [w(0, 0), w(0, 1)]
1371 /// [. , . , . , . ] (conv) [w(1, 0), w(1, 1)]
1372 /// [. , . , ., . ]
1373 /// [x(n-1, 0), x(n-1, 1), ..., x(n-1, n-1)]
1374 ///
1375 /// The packed input data (img2col) is a matrix with |rows| = output spatial
1376 /// size, |columns| = filter spatial size. To compute the output Y(i, j) we need
1377 /// to calculate the dot product between filter window at input X(x, y)) and the
1378 /// filter which will look like the following where r.h.s is the img2col matrix
1379 /// and l.h.s is the flattened filter:
1380 ///
1381 /// [x(0,0), x(0,1), x(1,0), x(1,1)]
1382 /// [x(0,1), x(1,1), x(0,2), x(1,2)] (matmul) [w(0,0), w(0,1), w(1,0), w(1,1)]
1383 /// [x(0,1), x(1,1), x(0,2), x(1,2)]
1384 /// [ . , . , . , . ]
1385 ///
1386 /// In general for 2D case with (N, H, W, C) input and (Kh, Kw, C, D) filter
1387 /// and output (N, Ho, Wo, D) the convolution is the following matrix-matrix
1388 /// multiplication (Ho x Wo, Kh x Kw x C) * (Kh x Kw x C, D) for each input in
1389 /// the N input. For the case where N > 1 its a batched matrix-matrix
1390 /// multiplication.
1391 ///
1392 /// On success, return both the operation that produces the img2col tensor and
1393 /// the final operation of the sequence that replaces the original convolution.
1394 FailureOr<std::pair<Operation *, Operation *>>
1395 rewriteInIm2Col(RewriterBase &rewriter, linalg::Conv2DNhwcHwcfOp convOp);
1396 
1397 /// Same as the above but for Fhwc channel orderings in the filter. In this case
1398 /// the matrix multiplication is actually a row-wise dot-product rather than a
1399 /// row-column dot-product. This is to avoid transposing the filter matrix which
1400 /// would be required for a regular matrix multiplication to produce the correct
1401 /// output dimensions.
1402 FailureOr<std::pair<Operation *, Operation *>>
1403 rewriteInIm2Col(RewriterBase &rewriter, linalg::Conv2DNhwcFhwcOp convOp);
1404 
1405 /// Similar to rewriteInIm2Col with linalg::Conv2DNhwcHwcfOp except there is no
1406 /// reduction among the input channels so each convolution can be a
1407 /// matrix-vector product and by transposing both input filter so channels are
1408 /// outer most the computation is a batched matrix-vector product.
1409 FailureOr<std::pair<Operation *, Operation *>>
1410 rewriteInIm2Col(RewriterBase &rewriter,
1411  linalg::DepthwiseConv2DNhwcHwcOp convOp);
1412 
1413 /// Similar to rewriteInIm2Col with linalg::Conv2DNhwcHwcfOp except because the
1414 /// channels are to the left of the image shape dimensions, the position of the
1415 /// contraction dimension in the resulting matmul is reversed. This swaps the
1416 /// LHS and RHS of the matmul when compared with nhwc (i.e. (D, C x Kh x Kw) *
1417 /// (C x Kh x Kw, Ho x Wo))
1418 FailureOr<std::pair<Operation *, Operation *>>
1419 rewriteInIm2Col(RewriterBase &rewriter, linalg::Conv2DNchwFchwOp convOp);
1420 
1421 /// Convert linalg.conv_2d_nhwc_fhwc(_q) to linalg.conv_2d_nhwc_hwcf(_q) by
1422 /// materializing transpose.
1423 FailureOr<Operation *> transposeConv2D(RewriterBase &rewriter,
1424  linalg::Conv2DNhwcFhwcOp op);
1425 FailureOr<Operation *> transposeConv2D(RewriterBase &rewriter,
1426  linalg::Conv2DNhwcFhwcQOp op);
1427 
1428 /// Convert Linalg matmul ops to transposed variants.
1429 FailureOr<Operation *> transposeMatmul(RewriterBase &rewriter,
1430  linalg::MatmulOp op,
1431  bool transposeLHS = true);
1432 FailureOr<Operation *> transposeBatchMatmul(RewriterBase &rewriter,
1433  linalg::BatchMatmulOp op,
1434  bool transposeLHS = true);
1435 
1436 /// Convert linalg.conv_2d_nhwc_fhwc to Winograd Conv2D algorithm
1437 /// F(m x m, r x r). m is the dimension size of output and r is the dimension
1438 /// size of filter.
1439 FailureOr<Operation *> winogradConv2D(RewriterBase &rewriter,
1440  linalg::Conv2DNhwcFhwcOp op,
1441  WinogradConv2DFmr fmr);
1442 
1443 /// Rewrite linalg.winograd_filter_transform. The data layout of the filter is
1444 /// FHWC. The transformation matrix is 2-dimension. We need to extract H x W
1445 /// from FHWC first. We generate 2 levels of loops to iterate on F and C. After
1446 /// the rewriting, we get
1447 ///
1448 /// scf.for %f = lo_f to hi_f step 1
1449 /// scf.for %c = lo_c to hi_c step 1
1450 /// %extracted = extract filter<h x w> from filter<f x h x w x c>
1451 /// %ret = linalg.matmul G, %extracted
1452 /// %ret = linalg.matmul %ret, GT
1453 /// %inserted = insert %ret into filter<h x w x c x f>
1454 FailureOr<Operation *>
1456  linalg::WinogradFilterTransformOp op);
1457 
1458 /// Rewrite linalg.winograd_input_transform. The data layout of the input is
1459 /// NHWC. The transformation matrix is 2-dimension. We need to extract H x W
1460 /// from NHWC first. We generate 4 levels of loops to iterate on N, C, tileH,
1461 /// and tileW. After the rewriting, we get
1462 ///
1463 /// scf.for %h = 0 to tileH step 1
1464 /// scf.for %w = 0 to tileW step 1
1465 /// scf.for %n = 0 to N step 1
1466 /// scf.for %c = 0 to C step 1
1467 /// %extracted = extract %extracted<alphaH x alphaW> from
1468 /// %input<N x H x W x C>
1469 /// at [%n, (%h x m), (%w x m), %c]
1470 /// %ret = linalg.matmul BT, %extracted
1471 /// %ret = linalg.matmul %ret, B
1472 /// %inserted = insert %ret<alphaH x alphaW> into
1473 /// %output<alphaH x alphaW x tileH x tileW x N x C>
1474 /// at [0, 0, %h, %w, %n, %c]
1475 FailureOr<Operation *>
1477  linalg::WinogradInputTransformOp op);
1478 
1479 /// Rewrite linalg.winograd_output_transform. The data layout of the output is
1480 /// HWNF. The transformation matrix is 2-dimension. We need to extract H x W
1481 /// from HWNF first. We generate 4 levels of loops to iterate on N, F, tileH,
1482 /// and tileW. After the transformation, we get
1483 ///
1484 /// scf.for %h = 0 to tileH step 1
1485 /// scf.for %w = 0 to tileW step 1
1486 /// scf.for %n = 0 to N step 1
1487 /// scf.for %f = 0 to F step 1
1488 /// %extracted = extract %extracted<alphaH x alphaW> from
1489 /// %input<alphaH x alphaW x tileH x tileW x N x F>
1490 /// at [0, 0, %h, %w, %n, %f]
1491 /// %ret = linalg.matmul AT, %extracted
1492 /// %ret = linalg.matmul %ret, A
1493 /// %inserted = insert %ret<alphaH x alphaW> into
1494 /// output<N x H x W x F>
1495 /// at [%n, (%h x m), (%w x m), %f]
1496 FailureOr<Operation *>
1498  linalg::WinogradOutputTransformOp op);
1499 
1500 /// Method to deduplicate operands and remove dead results of `linalg.generic`
1501 /// operations. This is effectively DCE for a linalg.generic op. If there is
1502 /// deduplication of operands orremoval of results, replaces the `genericOp`
1503 /// with a new op and returns it. Returns the same operation if there is no
1504 /// deduplication/removal.
1505 FailureOr<linalg::GenericOp> deduplicateOperandsAndRemoveDeadResults(
1506  RewriterBase &rewriter, linalg::GenericOp genericOp, bool removeOutputs);
1507 
1508 //===----------------------------------------------------------------------===//
1509 // Rewrite patterns wrapping transformations.
1510 // TODO: every single such pattern should be a close to noop wrapper around a
1511 // functional-stye API call.
1512 //===----------------------------------------------------------------------===//
1513 
1514 /// Rewrites 2-D convolution ops with size-1 window dimensions into 1-D
1515 /// convolution ops.
1516 template <typename Conv2DOp, typename Conv1DOp>
1518  : public OpRewritePattern<Conv2DOp> {
1520 
1521  FailureOr<Conv1DOp> returningMatchAndRewrite(Conv2DOp convOp,
1522  PatternRewriter &rewriter) const;
1523 
1524  LogicalResult matchAndRewrite(Conv2DOp convOp,
1525  PatternRewriter &rewriter) const override {
1526  return returningMatchAndRewrite(convOp, rewriter);
1527  }
1528 };
1529 
1530 extern template struct DownscaleSizeOneWindowed2DConvolution<Conv2DNhwcHwcfOp,
1531  Conv1DNwcWcfOp>;
1532 extern template struct DownscaleSizeOneWindowed2DConvolution<Conv2DNchwFchwOp,
1533  Conv1DNcwFcwOp>;
1534 
1535 /// Rewrites 2-D depthwise convolution ops with size-1 (w, kw) or (h, kh)
1536 /// dimensions into 1-D depthwise convolution ops.
1538  : public OpRewritePattern<DepthwiseConv2DNhwcHwcOp> {
1540  PatternBenefit benefit = 1)
1541  : OpRewritePattern<DepthwiseConv2DNhwcHwcOp>(context, benefit) {}
1542 
1543  FailureOr<DepthwiseConv1DNwcWcOp>
1544  returningMatchAndRewrite(DepthwiseConv2DNhwcHwcOp convOp,
1545  PatternRewriter &rewriter) const;
1546 
1547  LogicalResult matchAndRewrite(DepthwiseConv2DNhwcHwcOp convOp,
1548  PatternRewriter &rewriter) const override {
1549  return returningMatchAndRewrite(convOp, rewriter);
1550  }
1551 };
1552 
1553 struct DownscaleConv2DOp final : public OpRewritePattern<Conv2DOp> {
1555  : OpRewritePattern<Conv2DOp>(context, benefit) {}
1556 
1557  FailureOr<Conv1DOp> returningMatchAndRewrite(Conv2DOp convOp,
1558  PatternRewriter &rewriter) const;
1559 
1560  LogicalResult matchAndRewrite(Conv2DOp convOp,
1561  PatternRewriter &rewriter) const override {
1562  return returningMatchAndRewrite(convOp, rewriter);
1563  }
1564 };
1565 
1566 ///
1567 /// Linalg generalization pattern.
1568 ///
1569 /// Apply the `generalization` transformation as a pattern.
1570 /// See `generalization` for more details.
1571 //
1572 // TODO: Automatic default pattern class that just unwraps a function
1573 // returning FailureOr<GenericOp>.
1575  : public OpInterfaceRewritePattern<LinalgOp> {
1577 
1578  /// `matchAndRewrite` implementation that returns the significant
1579  /// transformed pieces of IR.
1580  FailureOr<GenericOp>
1581  returningMatchAndRewrite(LinalgOp op, PatternRewriter &rewriter) const {
1582  return generalizeNamedOp(rewriter, op);
1583  }
1584 
1585  LogicalResult matchAndRewrite(LinalgOp op,
1586  PatternRewriter &rewriter) const override {
1587  return returningMatchAndRewrite(op, rewriter);
1588  }
1589 };
1590 
1591 struct LinalgSpecializationPattern : public OpRewritePattern<GenericOp> {
1593 
1594  FailureOr<GenericOp>
1595  returningMatchAndRewrite(GenericOp op, PatternRewriter &rewriter) const {
1596  return specializeGenericOp(rewriter, op);
1597  }
1598 
1599  LogicalResult matchAndRewrite(GenericOp op,
1600  PatternRewriter &rewriter) const override {
1601  return returningMatchAndRewrite(op, rewriter);
1602  }
1603 };
1604 
1605 /// Vectorization pattern for memref::CopyOp.
1606 struct CopyVectorizationPattern : public OpRewritePattern<memref::CopyOp> {
1608 
1609  LogicalResult matchAndRewrite(memref::CopyOp copyOp,
1610  PatternRewriter &rewriter) const override;
1611 };
1612 
1614  std::function<LogicalResult(RewriterBase &, tensor::PadOp, Value)>;
1615 
1616 /// Rewrite a tensor::PadOp into a sequence of EmptyOp, FillOp and
1617 /// InsertSliceOp. For now, only constant padding values are supported.
1618 struct DecomposePadOpPattern : public OpRewritePattern<tensor::PadOp> {
1620  : OpRewritePattern<tensor::PadOp>(context, benefit) {}
1621  LogicalResult matchAndRewrite(tensor::PadOp padOp,
1622  PatternRewriter &rewriter) const override;
1623 
1624 protected:
1625  Value createFillOrGenerateOp(RewriterBase &rewriter, tensor::PadOp padOp,
1626  Value dest,
1627  const SmallVector<Value> &dynSizes) const;
1628 };
1629 
1630 /// Rewrites a linalg::PackOp into a sequence of:
1631 /// * tensor::PadOp + linalg::TransposeOp + tensor::EmptyOp +
1632 /// tensor::InsertSliceOp ops.
1633 ///
1634 /// Requires that all the outer dims of the input linalg::PackOp are 1.
1635 ///
1636 /// Before:
1637 /// ```
1638 /// %packed = linalg.pack %input
1639 /// padding_value(%pad : f32)
1640 /// inner_dims_pos = [1, 0]
1641 /// inner_tiles = [2, %high]
1642 /// into %output : tensor<5x1xf32> -> tensor<1x1x2x?xf32>
1643 /// ```
1644 ///
1645 /// After:
1646 /// ```
1647 /// // PadOp
1648 /// %padded = tensor.pad %arg0 low[0, 0] high[%0, 1] {
1649 /// ^bb0(...):
1650 /// tensor.yield %arg2 : f32
1651 /// } : tensor<5x1xf32> to tensor<?x2xf32>
1652 /// // EmptyOp + TransposeOp
1653 /// %empty = tensor.empty(%arg3) : tensor<2x?xf32>
1654 /// %transposed = linalg.transpose
1655 /// ins(%extracted_slice : tensor<?x2xf32>)
1656 /// outs(%empty : tensor<2x?xf32>)
1657 /// permutation = [1, 0]
1658 /// // InsertSliceOp
1659 /// %inserted_slice = tensor.insert_slice %transposed
1660 /// into %arg1[0, 0, 0, 0] [1, 1, 2, %tile_dim_1] [1, 1, 1, 1]
1661 /// : tensor<2x?xf32> into tensor<1x1x2x?xf32>
1662 /// ```
1664  : public OpRewritePattern<linalg::PackOp> {
1666  LogicalResult matchAndRewrite(linalg::PackOp packOp,
1667  PatternRewriter &rewriter) const override;
1668 };
1669 
1670 /// Rewrites a linalg::UnPackOp into a sequence of rank-reduced
1671 /// * tensor::ExtractSliceOp + linalg::TransposeOp + tensor::InsertSliceOp
1672 ///
1673 /// Requires that all the outer dims of the input linalg::PackOp are 1.
1674 ///
1675 /// Before:
1676 /// ```
1677 /// %packed = linalg.unpack %input
1678 /// inner_dims_pos = [1, 0]
1679 /// inner_tiles = [2, 8]
1680 /// into %output : tensor<1x1x2x8xf32> -> tensor<5x1xf32>
1681 /// ```
1682 ///
1683 /// After:
1684 /// ```
1685 /// // Rank-reduced extract to obtain the tile
1686 /// %slice = tensor.extract_slice %arg0[0, 0, 0, 0] [1, 1, 2, 8] [1, 1, 1, 1]
1687 /// : tensor<1x1x2x8xf32> to tensor<2x8xf32>
1688 /// // EmptyOp + TransposeOp
1689 /// %init = tensor.empty() : tensor<8x2xf32>
1690 /// %transposed = linalg.transpose
1691 /// ins(%extracted_slice : tensor<2x8xf32>)
1692 /// outs(%0 : tensor<8x2xf32>) permutation = [1, 0]
1693 /// // Extract a slice matching the specified output size
1694 /// %result = tensor.extract_slice %transposed[0, 0] [5, 1] [1, 1]
1695 /// : tensor<8x2xf32> to tensor<5x1xf32>
1696 /// ```
1698  : public OpRewritePattern<linalg::UnPackOp> {
1700  LogicalResult matchAndRewrite(linalg::UnPackOp unpackOp,
1701  PatternRewriter &rewriter) const override;
1702 };
1703 
1704 /// Match and rewrite for the pattern:
1705 /// ```
1706 /// %alloc = ...
1707 /// [optional] %view = memref.view %alloc ...
1708 /// %subView = subview %allocOrView ...
1709 /// [optional] linalg.fill(%allocOrView, %cst) ...
1710 /// ...
1711 /// memref.copy(%in, %subView) ...
1712 /// vector.transfer_read %allocOrView[...], %cst ...
1713 /// ```
1714 /// into
1715 /// ```
1716 /// [unchanged] %alloc = ...
1717 /// [unchanged] [optional] %view = memref.view %alloc ...
1718 /// [unchanged] [unchanged] %subView = subview %allocOrView ...
1719 /// ...
1720 /// vector.transfer_read %in[...], %cst ...
1721 /// ```
1722 /// Where there is no interleaved use between memref.copy and transfer_read as
1723 /// well as no interleaved use between linalg.fill and memref.copy (if
1724 /// linalg.fill is specified).
1725 /// This is a custom rewrite to forward partial reads (with optional fills) to
1726 /// vector.transfer_read.
1728  : public OpRewritePattern<vector::TransferReadOp> {
1730 
1731  LogicalResult matchAndRewrite(vector::TransferReadOp xferOp,
1732  PatternRewriter &rewriter) const override;
1733 };
1734 
1735 /// Match and rewrite for the pattern:
1736 /// ```
1737 /// %alloc = ...
1738 /// [optional] %view = memref.view %alloc ...
1739 /// %subView = subview %allocOrView...
1740 /// ...
1741 /// vector.transfer_write %..., %allocOrView[...]
1742 /// memref.copy(%subView, %out)
1743 /// ```
1744 /// into
1745 /// ```
1746 /// [unchanged] %alloc = ...
1747 /// [unchanged] [optional] %view = memref.view %alloc ...
1748 /// [unchanged] %subView = subview %allocOrView...
1749 /// ...
1750 /// vector.transfer_write %..., %out[...]
1751 /// ```
1752 /// Where there is no interleaved use between transfer_write and memref.copy.
1753 /// This is a custom rewrite to forward partial writes to
1754 /// vector.transfer_write.
1756  : public OpRewritePattern<vector::TransferWriteOp> {
1758 
1759  LogicalResult matchAndRewrite(vector::TransferWriteOp xferOp,
1760  PatternRewriter &rewriter) const override;
1761 };
1762 
1763 /// Rewrite extract_slice(tensor.pad(x)) into tensor.pad(extract_slice(x)).
1765  : public OpRewritePattern<tensor::ExtractSliceOp> {
1766  /// A function to control pattern application and rewrite logic.
1767  ///
1768  /// The function will be given the slice op and should return:
1769  /// - std::nullopt: to fail the match and not apply the pattern;
1770  /// - true: to apply the pattern with zero slice guard;
1771  /// - false: to apply the pattern without zero slice guard.
1772  ///
1773  /// See the documentation for tensor::bubbleUpPadSlice regarding zero slice
1774  /// guard.
1775  using ControlFn = std::function<std::optional<bool>(tensor::ExtractSliceOp)>;
1776 
1778  ControlFn controlFn = nullptr,
1779  PatternBenefit benefit = 1)
1780  : OpRewritePattern(context, benefit), controlFn(std::move(controlFn)) {}
1781 
1782  LogicalResult matchAndRewrite(tensor::ExtractSliceOp sliceOp,
1783  PatternRewriter &rewriter) const override;
1784 
1785 private:
1786  ControlFn controlFn;
1787 };
1788 
1789 //===----------------------------------------------------------------------===//
1790 // Populate functions.
1791 //===----------------------------------------------------------------------===//
1792 
1793 /// Canonicalization patterns relevant to apply after tiling patterns. These
1794 /// are applied automatically by the tiling pass but need to be applied
1795 /// manually when tiling is called programmatically.
1798 
1799 /// Linalg generalization patterns
1800 
1801 /// Populates `patterns` with patterns to convert spec-generated named ops to
1802 /// linalg.generic ops.
1804 
1805 /// Populates `patterns` with patterns to convert linalg.generic ops to named
1806 /// ops where possible. A linalg.generic can represent wide range and complex
1807 /// computations for which equivalent linalg named op may not exist e.g.
1808 /// linalg.generic that takes a tensor and computes a polynomial such as:
1809 /// p(x) = an*x^n + ... + a1x + a0
1810 /// There is no equivalent named op to convert to. Many such cases exist.
1813 
1814 /// Populates `patterns` with patterns that fold operations like
1815 /// `linalg.transform` into elementwise op map.
1817 
1818 /// Linalg decompose convolutions patterns
1819 
1820 /// Populates patterns to decompose high-D convolution ops into low-D ones.
1821 /// This is a step in progressive lowering for convolution ops, afterwards we
1822 /// can vectorize the low-D convolution ops.
1824  PatternBenefit benefit = 1);
1825 
1826 /// Populates patterns to decompose linalg.pack and linalg.unpack Ops into e.g.
1827 /// tensor.pad, linalg.transpose, tensor.{insert|extract}_slice. Require all
1828 /// outer dims to be unit.
1830 
1831 /// Populates patterns to decompose tensor.pad into e.g.
1832 /// tensor.empty, linalg.fill, tensor.insert_slice.
1834 
1835 /// Populates patterns to transform linalg.conv_2d_xxx operations into
1836 /// linalg.generic (for img2col packing) and linalg.matmul.
1837 /// \see rewriteInIm2Col for more details.
1839 
1840 /// Populates `patterns` with patterns that vectorize tensor.pad.
1841 /// These patterns are meant to apply in a complementary fashion. Benefits
1842 /// are used to encode a certain ordering of pattern application. To avoid
1843 /// scattering magic constants throughout the code base, the patterns must be
1844 /// added with this function. `baseBenefit` can be used to offset the benefit
1845 /// of all tensor::PadOp vectorization patterns by a certain value.
1847  PatternBenefit baseBenefit = 1);
1848 
1849 /// Populate patterns for splitting a `LinalgOp` with multiple statements within
1850 /// its payload into multiple `GenericOp` that have a single statement.
1851 /// The option `removeDeadArgsAndResults` adds patterns to remove dead arguments
1852 /// and results from the generated decomposed ops. This is default `true` since
1853 /// the core decomposition patterns relies on these clean up patterns. It is set
1854 /// to false only for testing purposes.
1856  bool removeDeadArgsAndResults = true);
1857 
1858 /// Populate patterns that convert non-destination-style ops to destination
1859 /// style ops.
1861 
1862 /// Populate patterns for vectorizing low-D convolution ops. This is a step in
1863 /// progressive lowering for convolution ops, it assume high-D convolution ops
1864 /// were decomposed previously.
1866  PatternBenefit benefit = 1);
1867 
1868 /// Populate patterns that convert `ElementwiseMappable` ops to linalg
1869 /// parallel loops.
1871 
1872 /// Populate patterns that are only useful in the context of sparse tensors.
1874 
1875 /// Function type which is used to control when to stop fusion. It is expected
1876 /// that OpOperand is not modified in the callback. The OpOperand is not marked
1877 /// as const to allow callers to use non-const methods.
1878 using ControlFusionFn = std::function<bool(OpOperand *fusedOperand)>;
1879 
1880 /// Patterns for fusing linalg operation on tensors.
1881 
1882 /// Pattern to fuse `linalg.generic` -> `linalg.generic` operations
1883 /// when both operations are fusable elementwise operations.
1886  const ControlFusionFn &controlElementwiseOpFusion);
1887 
1888 /// Function type which is used to control propagation of linalg.pack/unpack
1889 /// ops.
1890 using ControlPropagationFn = std::function<bool(OpOperand *opOperand)>;
1891 
1892 /// Patterns to bubble up or down data layout ops across other operations.
1895  const ControlPropagationFn &controlPackUnPackPropagation);
1896 
1897 /// Pattern to remove dead operands and results of `linalg.generic` operations.
1898 /// This is a pattern wrapper for `deduplicateOperandsAndRemoveDeadResults`.
1900 
1901 /// Patterns to promote inputs to outputs and remove unused inputs of
1902 /// `linalg.generic` ops.
1904 
1905 /// Function type to control generic op dimension collapsing. It is expected
1906 /// to return an array of `ReassociationIndices` representing dimensions that
1907 /// should be merged.
1909  std::function<SmallVector<ReassociationIndices>(linalg::LinalgOp)>;
1910 
1911 /// Pattern to collapse dimensions in a linalg.generic op. This will collapse
1912 /// tensor operands when needed and expand back the result tensors.
1915  const GetCollapsableDimensionsFn &controlCollapseDimensions);
1916 
1917 /// Patterns to fold an expanding (collapsing) tensor_reshape operation with its
1918 /// producer (consumer) generic operation by expanding the dimensionality of the
1919 /// loop in the generic op.
1921  RewritePatternSet &patterns, const ControlFusionFn &controlFoldingReshapes);
1922 
1923 /// Patterns to fold an expanding tensor.expand_shape operation with its
1924 /// producer generic operation by collapsing the dimensions of the generic op.
1926  RewritePatternSet &patterns, const ControlFusionFn &controlFoldingReshapes);
1927 
1928 /// Patterns to constant fold Linalg operations.
1930  const ControlFusionFn &controlFn);
1931 
1932 /// Pattern to replace `linalg.add` when destination passing on a contraction op
1933 /// suffices for achieving the sum.
1935 
1936 /// Pattern to fuse a `tensor.pad` operation with the producer of its source,
1937 /// if the producer is a `linalg` operation with all parallel iterator types.
1940 
1941 /// Patterns to convert from one named op to another. These can be seen as
1942 /// canonicalizations of named ops into another named op.
1944 
1945 /// Patterns to fold unit-extent dimensions in operands/results of linalg ops on
1946 /// tensors via reassociative reshape ops.
1949 
1950 /// A pattern that converts init operands to input operands.
1952 
1953 /// Patterns that are used to inline constant operands into linalg generic ops.
1955 
1956 /// Patterns that are used to bubble up extract slice op above linalg op.
1958 
1959 /// Adds patterns that waps tensor.extract_slice(linalg.fill(%cst, %init)) into
1960 /// linalg.fill(%cst, tensor.extract_slice(%init)).
1962 
1963 /// Add patterns to make explicit broadcasts and transforms in the
1964 /// input operands of a genericOp.
1966 
1967 /// Patterns to apply `splitReduction` below.
1970  const ControlSplitReductionFn &controlSplitReductionFn,
1971  bool useAlloc = false);
1972 
1973 /// Patterns to convert Linalg matmul ops to transposed variants.
1975  bool transposeLHS = true);
1976 
1977 /// Patterns to block pack Linalg matmul ops.
1979  const ControlBlockPackMatmulFn &controlFn);
1980 
1981 /// Patterns to apply Winograd Conv2D algorithm F(m x m, r x r).
1983  WinogradConv2DFmr fmr);
1984 
1985 /// Patterns to decompose Winograd operators.
1987 
1988 /// Adds patterns that reduce the rank of named contraction ops that have
1989 /// unit dimensions in the operand(s) by converting to a sequence of
1990 /// `collapse_shape`,
1991 /// `<corresponding linalg named op>`, `expand_shape` (if on tensors). For
1992 /// example a `linalg.batch_matmul` with unit batch size will convert to
1993 /// `linalg.matmul` and a `linalg.matvec` with with unit spatial dim in lhs will
1994 /// convert to a `linalg.dot`.
1996 
1997 /// Function type which is used to control folding operations like `tensor.pad`
1998 /// and `tensor.extract_slice` into linalg.pack/unpack ops.
1999 using ControlFoldIntoPackUnpackFn = std::function<bool(OpOperand *opOperand)>;
2000 /// Populates `patterns` with patterns that fold operations like `tensor.pad`
2001 /// and `tensor.extract_slice` into `tensor.pack` and `tensor.unpack` operations
2002 /// respectively.
2005  const ControlFoldIntoPackUnpackFn &controlFn = nullptr);
2006 
2007 /// Populates `patterns` with patterns that fold operations like `linalg.pack`
2008 /// and `linalg.unpack` into `tensor.empty`.
2010 
2011 /// Populates `patterns` with patterns that simplify `tensor.pack` and
2012 /// `tensor.unpack` operations.
2014 
2015 } // namespace linalg
2016 } // namespace mlir
2017 
2018 #endif // MLIR_DIALECT_LINALG_TRANSFORMS_TRANSFORMS_H
static llvm::ManagedStatic< PassManagerOptions > options
A multi-dimensional affine map Affine map's are immutable like Type's, and they are uniqued.
Definition: AffineMap.h:46
Attributes are known-constant values of operations.
Definition: Attributes.h:25
The main mechanism for performing data layout queries.
This class defines the main interface for locations in MLIR and acts as a non-nullable wrapper around...
Definition: Location.h:76
MLIRContext is the top-level object for a collection of MLIR operations.
Definition: MLIRContext.h:60
This class helps build Operations.
Definition: Builders.h:205
This class represents a single result from folding an operation.
Definition: OpDefinition.h:272
This class represents an operand of an operation.
Definition: Value.h:257
Operation is the basic unit of execution within MLIR.
Definition: Operation.h:88
This class represents the benefit of a pattern match in a unitless scheme that ranges from 0 (very li...
Definition: PatternMatch.h:34
A special type of RewriterBase that coordinates the application of a rewrite pattern on the current I...
Definition: PatternMatch.h:767
This class coordinates the application of a rewrite on a set of IR, providing a way for clients to tr...
Definition: PatternMatch.h:358
This class represents an instance of an SSA value in the MLIR system, representing a computable value...
Definition: Value.h:96
FailureOr< PackingResult > buildPackingLoopNest(RewriterBase &rewriter, tensor::PadOp opToHoist, scf::ForOp outermostEnclosingForOp, ArrayRef< int64_t > transposeVector)
Build the packing loop nest required to hoist opToHoist above outermostEnclosingForOp.
void populateLinalgNamedOpConversionPatterns(RewritePatternSet &patterns)
Patterns to convert from one named op to another.
void populateMoveInitOperandsToInputPattern(RewritePatternSet &patterns)
A pattern that converts init operands to input operands.
void populateTransposeMatmulPatterns(RewritePatternSet &patterns, bool transposeLHS=true)
Patterns to convert Linalg matmul ops to transposed variants.
void populateContractionOpRankReducingPatterns(RewritePatternSet &patterns)
Adds patterns that reduce the rank of named contraction ops that have unit dimensions in the operand(...
LogicalResult rewriteAsPaddedOp(RewriterBase &rewriter, LinalgOp opToPad, const LinalgPaddingOptions &options, LinalgOp &paddedOp, SmallVector< Value > &replacements, SmallVector< tensor::PadOp > &padOps)
Pad the iterator dimensions options.paddingDimensions of all opToPad operands to a static bounding bo...
Definition: Padding.cpp:244
void populateSplitReductionPattern(RewritePatternSet &patterns, const ControlSplitReductionFn &controlSplitReductionFn, bool useAlloc=false)
Patterns to apply splitReduction below.
void populateFuseTensorPadWithProducerLinalgOpPatterns(RewritePatternSet &patterns)
Pattern to fuse a tensor.pad operation with the producer of its source, if the producer is a linalg o...
FailureOr< std::pair< Operation *, Operation * > > rewriteInIm2Col(RewriterBase &rewriter, linalg::Conv2DNhwcHwcfOp convOp)
Convert linalg.conv_2d_nhwc_hwcf into linalg.generic (for img2col packing) and linalg....
bool areDimSequencesPreserved(ArrayRef< AffineMap > maps, ArrayRef< ReassociationIndices > dimSequences)
Return true if all sequences of dimensions specified in dimSequences are contiguous in all the ranges...
bool hasVectorizationImpl(Operation *)
Return true if there's dedicated logic in the Linalg Vectorizer to vectorize this Op,...
void populateBubbleUpExtractSliceOpPatterns(RewritePatternSet &patterns)
Patterns that are used to bubble up extract slice op above linalg op.
void transformIndexOps(RewriterBase &b, LinalgOp op, SmallVectorImpl< Value > &ivs, const LoopIndexToRangeIndexMap &loopIndexToRangeIndex)
All indices returned by IndexOp should be invariant with respect to tiling.
Definition: Tiling.cpp:73
std::function< std::optional< Value >(OpBuilder &b, memref::SubViewOp subView, ArrayRef< Value > boundingSubViewSize, DataLayout &layout)> AllocBufferCallbackFn
Callback function type used to perform the allocation for the promoted subView.
Definition: Transforms.h:381
void populateBlockPackMatmulPatterns(RewritePatternSet &patterns, const ControlBlockPackMatmulFn &controlFn)
Patterns to block pack Linalg matmul ops.
void populateConvertConv2DToImg2ColPatterns(RewritePatternSet &patterns)
Populates patterns to transform linalg.conv_2d_xxx operations into linalg.generic (for img2col packin...
FailureOr< Operation * > decomposeWinogradFilterTransformOp(RewriterBase &rewriter, linalg::WinogradFilterTransformOp op)
Rewrite linalg.winograd_filter_transform.
DenseMap< int, int > LoopIndexToRangeIndexMap
Creates a number of ranges equal to the number of non-zero in tileSizes.
Definition: Transforms.h:910
std::optional< Value > allocateWorkgroupMemory(OpBuilder &builder, memref::SubViewOp subview, ArrayRef< Value > sizeBounds, DataLayout &)
Allocate the subview in the GPU workgroup memory.
Definition: Promotion.cpp:472
FailureOr< PackTransposeResult > packTranspose(RewriterBase &rewriter, linalg::PackOp packOp, linalg::LinalgOp linalgOp, linalg::UnPackOp maybeUnPackOp, ArrayRef< int64_t > outerPerm, ArrayRef< int64_t > innerPerm)
Transpose a single PackOp -> LinalgOp -> UnPackOp chain and return the transposed PackOp -> LinalgOp ...
Definition: Transforms.cpp:673
Value bufferizeToAllocation(RewriterBase &rewriter, const BufferizeToAllocationOptions &options, tensor::PadOp padOp, Attribute memorySpace={}, Operation *insertionPoint=nullptr)
Materialize a buffer allocation for the given tensor.pad op and lower the op to linalg....
std::function< bool(OpOperand *fusedOperand)> ControlFusionFn
Function type which is used to control when to stop fusion.
Definition: Transforms.h:1878
bool isDimSequencePreserved(AffineMap map, ReassociationIndicesRef dimSequence)
Return true if a given sequence of dimensions are contiguous in the range of the specified indexing m...
FailureOr< Value > hoistPaddingOnTensors(RewriterBase &rewriter, tensor::PadOp opToHoist, int64_t numLoops, ArrayRef< int64_t > transposeVector, tensor::PadOp &hoistedOp, SmallVectorImpl< TransposeOp > &transposeOps)
Mechanically hoist padding operations on tensors by numLoops into a new, generally larger tensor.
SmallVector< OpFoldResult > computePaddedShape(RewriterBase &rewriter, TypedValue< RankedTensorType > v, AffineMap indexingMap, ArrayRef< OpFoldResult > indexingSizes, const PadTilingInterfaceOptions &options)
Helper function to compute the padded shape of the given value v of RankedTensorType given:
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.
Definition: Specialize.cpp:258
void populateFoldReshapeOpsByCollapsingPatterns(RewritePatternSet &patterns, const ControlFusionFn &controlFoldingReshapes)
Patterns to fold an expanding tensor.expand_shape operation with its producer generic operation by co...
LinalgTilingLoopType
The type of loops to be generated during tiling.
Definition: Utils.h:118
FailureOr< LowerUnPackOpResult > lowerUnPack(RewriterBase &rewriter, linalg::UnPackOp unPackOp, bool lowerUnpadLikeWithExtractSlice=true)
Rewrite pack as empty + transpose + reshape + extract_slice.
Definition: Transforms.cpp:354
std::function< LogicalResult(OpBuilder &b, Value buffer)> DeallocBufferCallbackFn
Callback function type used to deallocate the buffers used to hold the promoted subview.
Definition: Transforms.h:386
void populateDataLayoutPropagationPatterns(RewritePatternSet &patterns, const ControlPropagationFn &controlPackUnPackPropagation)
Patterns to bubble up or down data layout ops across other operations.
void populatePadOpVectorizationPatterns(RewritePatternSet &patterns, PatternBenefit baseBenefit=1)
Populates patterns with patterns that vectorize tensor.pad.
void populateLinalgTilingCanonicalizationPatterns(RewritePatternSet &patterns)
Definition: Tiling.cpp:856
void populateLinalgFoldIntoElementwisePatterns(RewritePatternSet &patterns)
Populates patterns with patterns that fold operations like linalg.transform into elementwise op map.
LogicalResult deallocateGPUPrivateMemory(OpBuilder &, Value)
In case of GPU private memory there is no need to deallocate since the memory is freed when going out...
Definition: Promotion.cpp:513
void populateSparseTensorRewriting(RewritePatternSet &patterns)
Populate patterns that are only useful in the context of sparse tensors.
FailureOr< Operation * > decomposeWinogradOutputTransformOp(RewriterBase &rewriter, linalg::WinogradOutputTransformOp op)
Rewrite linalg.winograd_output_transform.
void populateWinogradConv2DPatterns(RewritePatternSet &patterns, WinogradConv2DFmr fmr)
Patterns to apply Winograd Conv2D algorithm F(m x m, r x r).
FailureOr< ElementwiseOpFusionResult > fuseElementwiseOps(RewriterBase &rewriter, OpOperand *fusedOperand)
FailureOr< PromotionInfo > promoteSubviewAsNewBuffer(OpBuilder &b, Location loc, memref::SubViewOp subView, bool useOriginalSubviewSize, const AllocBufferCallbackFn &allocationFn, DataLayout &layout)
Definition: Promotion.cpp:238
llvm::SmallDenseSet< int > getPreservedProducerResults(GenericOp producer, GenericOp consumer, OpOperand *fusedOperand)
Returns a set of indices of the producer's results which would be preserved after the fusion.
std::optional< Value > allocateGPUPrivateMemory(OpBuilder &builder, memref::SubViewOp subview, ArrayRef< Value > sizeBounds, DataLayout &)
Allocate the subview in the GPU private memory.
Definition: Promotion.cpp:497
FailureOr< Operation * > rewriteInDestinationPassingStyle(RewriterBase &rewriter, tensor::FromElementsOp fromElementsOp)
Rewrite tensor.from_elements to linalg.generic.
FailureOr< DropUnitDimsResult > dropUnitDims(RewriterBase &rewriter, GenericOp genericOp, const ControlDropUnitDims &options)
FailureOr< PackResult > blockPackMatmul(RewriterBase &rewriter, linalg::LinalgOp linalgOp, const ControlBlockPackMatmulFn &controlPackMatmul)
Pack a matmul operation into blocked 4D layout.
void peelLoops(RewriterBase &rewriter, ArrayRef< scf::ForOp > loops)
Peel 'loops' and applies affine_min/max bounds simplification on the fly where relevant.
Definition: Transforms.cpp:71
void populateConvertToDestinationStylePatterns(RewritePatternSet &patterns)
Populate patterns that convert non-destination-style ops to destination style ops.
FailureOr< Operation * > transposeConv2D(RewriterBase &rewriter, linalg::Conv2DNhwcFhwcOp op)
Convert linalg.conv_2d_nhwc_fhwc(_q) to linalg.conv_2d_nhwc_hwcf(_q) by materializing transpose.
void populateFoldUnitExtentDimsPatterns(RewritePatternSet &patterns, ControlDropUnitDims &options)
Patterns to fold unit-extent dimensions in operands/results of linalg ops on tensors via reassociativ...
LogicalResult copyToWorkgroupMemory(OpBuilder &b, Value src, Value dst)
Create Memref copy operations and add gpu barrier guards before and after the copy operation to ensur...
Definition: Promotion.cpp:488
std::function< SmallVector< Value, 4 >(OpBuilder &, Operation *)> TileSizeComputationFunction
Definition: Transforms.h:189
std::function< LogicalResult(RewriterBase &, tensor::PadOp, Value)> OptimizeCopyFn
Definition: Transforms.h:1614
void populateElementwiseToLinalgConversionPatterns(RewritePatternSet &patterns)
Populate patterns that convert ElementwiseMappable ops to linalg parallel loops.
LogicalResult linalgOpAnchoredEmptyTensorEliminationStep(RewriterBase &rewriter, Operation *op, bufferization::OneShotAnalysisState &state)
Try to eliminate tensor::EmptyOps inside op that are anchored on a LinalgOp.
FailureOr< LinalgLoops > linalgOpToLoops(RewriterBase &rewriter, LinalgOp linalgOp)
Emit a loop nest of scf.for with the proper body for linalgOp.
Definition: Loops.cpp:367
FailureOr< GenericOp > generalizeNamedOp(RewriterBase &rewriter, LinalgOp linalgOp)
Create a GenericOp from the given named operation linalgOp and replace the given linalgOp.
std::tuple< SmallVector< Range, 4 >, LoopIndexToRangeIndexMap > makeTiledLoopRanges(RewriterBase &b, Location loc, AffineMap map, ArrayRef< OpFoldResult > allShapeSizes, ArrayRef< OpFoldResult > allTileSizes)
Definition: Tiling.cpp:44
FailureOr< Operation * > transposeBatchMatmul(RewriterBase &rewriter, linalg::BatchMatmulOp op, bool transposeLHS=true)
Pattern to replace.
LogicalResult promoteSubviewsPrecondition(Operation *op, LinalgPromotionOptions options)
Promote memref.subviews feeding linalg-on-buffers operations.
Definition: Promotion.cpp:401
LogicalResult copyToGPUPrivateMemory(OpBuilder &b, Value src, Value dst)
Normal copy to between src and dst.
Definition: Promotion.cpp:505
FailureOr< linalg::GenericOp > deduplicateOperandsAndRemoveDeadResults(RewriterBase &rewriter, linalg::GenericOp genericOp, bool removeOutputs)
Method to deduplicate operands and remove dead results of linalg.generic operations.
void populateDecomposeConvolutionPatterns(RewritePatternSet &patterns, PatternBenefit benefit=1)
Linalg decompose convolutions patterns.
void populateDecomposeWinogradOpsPatterns(RewritePatternSet &patterns)
Patterns to decompose Winograd operators.
void populateConvolutionVectorizationPatterns(RewritePatternSet &patterns, PatternBenefit benefit=1)
Populate patterns for vectorizing low-D convolution ops.
std::function< bool(OpOperand *opOperand)> ControlFoldIntoPackUnpackFn
Function type which is used to control folding operations like tensor.pad and tensor....
Definition: Transforms.h:1999
FailureOr< Operation * > winogradConv2D(RewriterBase &rewriter, linalg::Conv2DNhwcFhwcOp op, WinogradConv2DFmr fmr)
Convert linalg.conv_2d_nhwc_fhwc to Winograd Conv2D algorithm F(m x m, r x r).
LogicalResult vectorizeCopy(RewriterBase &builder, memref::CopyOp copyOp)
Emit a suitable vector form for a Copy op with fully static shape.
FailureOr< SmallVector< OpFoldResult > > computeIndexingMapOpInterfacePaddedShape(RewriterBase &rewriter, OpOperand &operandToPad, ArrayRef< Range > iterationDomain, const PadTilingInterfaceOptions &options)
Specific helper for Linalg ops.
LogicalResult vectorizeOpPrecondition(Operation *op, ArrayRef< int64_t > inputVectorSizes={}, ArrayRef< bool > inputScalableVecDims={}, bool vectorizeNDExtract=false, bool flatten1DDepthwiseConv=false)
Return success if the operation can be vectorized.
FailureOr< GenericOp > interchangeGenericOp(RewriterBase &rewriter, GenericOp genericOp, ArrayRef< unsigned > interchangeVector)
Interchange the iterator_types and iterator_maps dimensions and adapts the index accesses of op.
Definition: Interchange.cpp:45
void populateCollapseDimensions(RewritePatternSet &patterns, const GetCollapsableDimensionsFn &controlCollapseDimensions)
Pattern to collapse dimensions in a linalg.generic op.
bool areElementwiseOpsFusable(OpOperand *fusedOperand)
Return true if two linalg.generic operations with producer/consumer relationship through fusedOperand...
FailureOr< StaticMultiSizeSpecification > computeStaticMultiTileSizes(LinalgOp op, unsigned dimension, int64_t targetSize, int64_t divisor)
Definition: Tiling.cpp:236
FailureOr< LinalgLoops > linalgOpToAffineLoops(RewriterBase &rewriter, LinalgOp linalgOp)
Emit a loop nest of affine.for with the proper body for linalgOp.
Definition: Loops.cpp:362
void populateDecomposePackUnpackPatterns(RewritePatternSet &patterns)
Populates patterns to decompose linalg.pack and linalg.unpack Ops into e.g.
void populateEraseUnusedOperandsAndResultsPatterns(RewritePatternSet &patterns)
Pattern to remove dead operands and results of linalg.generic operations.
FailureOr< ContinuousTileSizeSpecification > computeContinuousTileSizes(OpBuilder &builder, TilingInterface op, unsigned dimension, OpFoldResult targetSize, bool emitAssertions)
Definition: Tiling.cpp:156
FailureOr< StaticContinuousTileSizeSpecification > computeStaticContinuousTileSizes(LinalgOp op, unsigned dimension, unsigned targetSize)
Definition: Tiling.cpp:106
std::function< LogicalResult(OpBuilder &b, Value src, Value dst)> CopyCallbackFn
Callback function type used to insert copy from original subview to subview of the promoted region fo...
Definition: Transforms.h:393
FailureOr< SplitReductionResult > splitReduction(RewriterBase &b, LinalgOp op, const ControlSplitReductionFn &controlSplitReductionFn, bool useAlloc=false)
void populateSimplifyPackAndUnpackPatterns(RewritePatternSet &patterns)
Populates patterns with patterns that simplify tensor.pack and tensor.unpack operations.
void populateFoldPackUnpackIntoTensorEmptyPatterns(RewritePatternSet &patterns)
Populates patterns with patterns that fold operations like linalg.pack and linalg....
FailureOr< LinalgOp > padAndHoistLinalgOp(RewriterBase &rewriter, LinalgOp linalgOp, const LinalgPaddingOptions &options)
Apply padding and hoisting to linalgOp according to the configuration specified in options.
Definition: Padding.cpp:356
void populateDecomposeLinalgOpsPattern(RewritePatternSet &patterns, bool removeDeadArgsAndResults=true)
Populate patterns for splitting a LinalgOp with multiple statements within its payload into multiple ...
std::function< bool(OpOperand *opOperand)> ControlPropagationFn
Function type which is used to control propagation of linalg.pack/unpack ops.
Definition: Transforms.h:1890
void populateFoldIntoPackAndUnpackPatterns(RewritePatternSet &patterns, const ControlFoldIntoPackUnpackFn &controlFn=nullptr)
Populates patterns with patterns that fold operations like tensor.pad and tensor.extract_slice into t...
FailureOr< ForallReductionTilingResult > tileReductionUsingForall(RewriterBase &b, PartialReductionOpInterface op, ArrayRef< OpFoldResult > numThreads, ArrayRef< OpFoldResult > tileSizes={}, std::optional< ArrayAttr > mapping=std::nullopt)
Method to tile a reduction to parallel iterations computing partial reductions.
Definition: Tiling.cpp:588
FailureOr< PackResult > packMatmulGreedily(RewriterBase &rewriter, LinalgOp linalgOp, ArrayRef< OpFoldResult > mnkPackedSizes, ArrayRef< int64_t > mnkPaddedSizesNextMultipleOf, ArrayRef< int64_t > mnkOrder)
Pack a LinalgOp by greedily inferring matmul dimensions (m, n, k) where m and n are proper parallel d...
Definition: Transforms.cpp:764
FailureOr< PackResult > pack(RewriterBase &rewriter, linalg::LinalgOp linalgOp, ArrayRef< OpFoldResult > packedSizes)
Implement packing of a single LinalgOp by packedSizes.
Definition: Transforms.cpp:476
void populateEraseUnnecessaryInputsPatterns(RewritePatternSet &patterns)
Patterns to promote inputs to outputs and remove unused inputs of linalg.generic ops.
FailureOr< TiledLinalgOp > tileLinalgOp(RewriterBase &b, LinalgOp op, const LinalgTilingOptions &options)
Definition: Tiling.cpp:816
std::function< SmallVector< ReassociationIndices >(linalg::LinalgOp)> GetCollapsableDimensionsFn
Function type to control generic op dimension collapsing.
Definition: Transforms.h:1909
std::function< FailureOr< SmallVector< OpFoldResult > >(RewriterBase &, OpOperand &, ArrayRef< Range >, const PadTilingInterfaceOptions &)> PadSizeComputationFunction
Definition: Transforms.h:614
void populateFoldReshapeOpsByExpansionPatterns(RewritePatternSet &patterns, const ControlFusionFn &controlFoldingReshapes)
Patterns to fold an expanding (collapsing) tensor_reshape operation with its producer (consumer) gene...
void populateSwapExtractSliceWithFillPatterns(RewritePatternSet &patterns)
Adds patterns that waps tensor.extract_slice(linalg.fill(cst, init)) into linalg.fill(cst,...
void populateInlineConstantOperandsPatterns(RewritePatternSet &patterns)
Patterns that are used to inline constant operands into linalg generic ops.
FailureOr< LinalgOp > promoteSubViews(OpBuilder &b, LinalgOp op, const LinalgPromotionOptions &options)
Promote the subViews into a new buffer allocated at the insertion point b.
Definition: Promotion.cpp:423
void populateConstantFoldLinalgOperations(RewritePatternSet &patterns, const ControlFusionFn &controlFn)
Patterns to constant fold Linalg operations.
std::function< SplitReductionOptions(LinalgOp op)> ControlSplitReductionFn
Function signature to control reduction splitting.
Definition: Transforms.h:491
LogicalResult deallocateWorkgroupMemory(OpBuilder &, Value)
In case of GPU group memory there is no need to deallocate.
Definition: Promotion.cpp:481
FailureOr< Operation * > transposeMatmul(RewriterBase &rewriter, linalg::MatmulOp op, bool transposeLHS=true)
Convert Linalg matmul ops to transposed variants.
FailureOr< VectorizationResult > vectorize(RewriterBase &rewriter, Operation *op, ArrayRef< int64_t > inputVectorSizes={}, ArrayRef< bool > inputScalableVecDims={}, bool vectorizeNDExtract=false, bool flatten1DDepthwiseConv=false, bool assumeDynamicDimsMatchVecSizes=false)
Returns a VectorizationResult containing the results of the vectorized op, or failure if the transfor...
void populateLinalgNamedOpsGeneralizationPatterns(RewritePatternSet &patterns)
Linalg generalization patterns.
void populateLinalgGenericOpsSpecializationPatterns(RewritePatternSet &patterns)
Populates patterns with patterns to convert linalg.generic ops to named ops where possible.
Definition: Specialize.cpp:355
std::function< std::optional< BlockPackMatmulOptions >(linalg::LinalgOp)> ControlBlockPackMatmulFn
Function type which is used to control matmul packing.
Definition: Transforms.h:1319
enum WinogradConv2DFmr uint32_t std::optional< vector::CombiningKind > getCombinerOpKind(Operation *combinerOp)
Return vector::CombiningKind for the given op.
SmallVector< Value > peelLoop(RewriterBase &rewriter, Operation *op)
Try to peel and canonicalize loop op and return the new result.
Definition: Transforms.cpp:55
RewritePatternSet getLinalgTilingCanonicalizationPatterns(MLIRContext *ctx)
Canonicalization patterns relevant to apply after tiling patterns.
Definition: Tiling.cpp:850
FailureOr< CollapseResult > collapseOpIterationDims(LinalgOp op, ArrayRef< ReassociationIndices > foldedIterationDims, RewriterBase &rewriter)
Collapses dimensions of linalg.generic/linalg.copy operation.
FailureOr< Operation * > decomposeWinogradInputTransformOp(RewriterBase &rewriter, linalg::WinogradInputTransformOp op)
Rewrite linalg.winograd_input_transform.
void populateDecomposePadPatterns(RewritePatternSet &patterns)
Populates patterns to decompose tensor.pad into e.g.
void populateFoldAddIntoDestPatterns(RewritePatternSet &patterns)
Pattern to replace linalg.add when destination passing on a contraction op suffices for achieving the...
std::pair< TilingInterface, TilingInterface > splitOp(RewriterBase &rewriter, TilingInterface op, unsigned dimension, OpFoldResult splitPoint)
Split the given op into two parts along the given iteration space dimension at the specified splitPoi...
Definition: Split.cpp:67
void populateElementwiseOpsFusionPatterns(RewritePatternSet &patterns, const ControlFusionFn &controlElementwiseOpFusion)
Patterns for fusing linalg operation on tensors.
FailureOr< SplitReductionResult > splitReductionByScaling(RewriterBase &b, LinalgOp op, const ControlSplitReductionFn &controlSplitReductionFn, bool useAlloc=false)
Scaling-based implementation of the split reduction transformation.
FailureOr< MultiSizeSpecification > computeMultiTileSizes(OpBuilder &builder, LinalgOp op, unsigned dimension, OpFoldResult targetSize, OpFoldResult divisor, bool emitAssertions=true)
Emits the IR computing the multi-sized tiling specification with two tile sizes not exceeding targetS...
Definition: Tiling.cpp:262
FailureOr< LowerPackResult > lowerPack(RewriterBase &rewriter, linalg::PackOp packOp, bool lowerPadLikeWithInsertSlice=true)
Rewrite pack as pad + reshape + transpose.
Definition: Transforms.cpp:219
FailureOr< LinalgLoops > linalgOpToParallelLoops(RewriterBase &rewriter, LinalgOp linalgOp)
Emit a loop nest of scf.parallel with the proper body for linalgOp.
Definition: Loops.cpp:374
Include the generated interface declarations.
std::conditional_t< std::is_same_v< Ty, mlir::Type >, mlir::Value, detail::TypedValue< Ty > > TypedValue
If Ty is mlir::Type this will select Value instead of having a wrapper around it.
Definition: Value.h:488
ArrayRef< int64_t > ReassociationIndicesRef
const FrozenRewritePatternSet & patterns
OpInterfaceRewritePattern is a wrapper around RewritePattern that allows for matching and rewriting a...
Definition: PatternMatch.h:330
OpRewritePattern is a wrapper around RewritePattern that allows for matching and rewriting against an...
Definition: PatternMatch.h:314
SmallVector< int64_t, 3 > mnkOrder
Permutation of matmul (M, N, K) dimensions order.
Definition: Transforms.h:1299
SmallVector< int64_t, 3 > blockFactors
Minor block factors (mb, nb, kb) for packing relayout where mb, mn are the parallel dimensions and kb...
Definition: Transforms.h:1289
bool rhsTransposeOuterBlocks
Transpose RHS outer block layout [KB][NB] -> [NB][KB].
Definition: Transforms.h:1308
bool lhsTransposeInnerBlocks
Transpose LHS inner block layout [mb][kb] -> [kb][mb].
Definition: Transforms.h:1305
SmallVector< int64_t, 3 > mnkPaddedSizesNextMultipleOf
Next multiples of the packing sizes.
Definition: Transforms.h:1296
bool lhsTransposeOuterBlocks
Transpose LHS outer block layout [MB][KB] -> [KB][MB].
Definition: Transforms.h:1302
bool allowPadding
If true, allows packing of dimensions that only partially fit into the block factors.
Definition: Transforms.h:1293
bool rhsTransposeInnerBlocks
Transpose RHS inner block layout [kb][nb] -> [nb][kb].
Definition: Transforms.h:1311
SmallVector< Value > results
Definition: Transforms.h:1202
Transformation to drop unit-extent dimensions from linalg.generic operations.
Definition: Transforms.h:522
RankReductionStrategy rankReductionStrategy
Definition: Transforms.h:525
std::function< SmallVector< unsigned >(Operation *)> ControlFnTy
Definition: Transforms.h:528
Vectorization pattern for memref::CopyOp.
Definition: Transforms.h:1606
LogicalResult matchAndRewrite(memref::CopyOp copyOp, PatternRewriter &rewriter) const override
Definition: Transforms.cpp:914
Rewrites a linalg::PackOp into a sequence of:
Definition: Transforms.h:1664
LogicalResult matchAndRewrite(linalg::PackOp packOp, PatternRewriter &rewriter) const override
Rewrites a linalg::UnPackOp into a sequence of rank-reduced.
Definition: Transforms.h:1698
LogicalResult matchAndRewrite(linalg::UnPackOp unpackOp, PatternRewriter &rewriter) const override
Rewrite a tensor::PadOp into a sequence of EmptyOp, FillOp and InsertSliceOp.
Definition: Transforms.h:1618
LogicalResult matchAndRewrite(tensor::PadOp padOp, PatternRewriter &rewriter) const override
Definition: Transforms.cpp:942
Value createFillOrGenerateOp(RewriterBase &rewriter, tensor::PadOp padOp, Value dest, const SmallVector< Value > &dynSizes) const
Filling dest using FillOp constant padding value if possible.
Definition: Transforms.cpp:921
DecomposePadOpPattern(MLIRContext *context, PatternBenefit benefit=1)
Definition: Transforms.h:1619
LogicalResult matchAndRewrite(Conv2DOp convOp, PatternRewriter &rewriter) const override
Definition: Transforms.h:1560
FailureOr< Conv1DOp > returningMatchAndRewrite(Conv2DOp convOp, PatternRewriter &rewriter) const
DownscaleConv2DOp(MLIRContext *context, PatternBenefit benefit=1)
Definition: Transforms.h:1554
Rewrites 2-D depthwise convolution ops with size-1 (w, kw) or (h, kh) dimensions into 1-D depthwise c...
Definition: Transforms.h:1538
FailureOr< DepthwiseConv1DNwcWcOp > returningMatchAndRewrite(DepthwiseConv2DNhwcHwcOp convOp, PatternRewriter &rewriter) const
LogicalResult matchAndRewrite(DepthwiseConv2DNhwcHwcOp convOp, PatternRewriter &rewriter) const override
Definition: Transforms.h:1547
DownscaleDepthwiseConv2DNhwcHwcOp(MLIRContext *context, PatternBenefit benefit=1)
Definition: Transforms.h:1539
Rewrites 2-D convolution ops with size-1 window dimensions into 1-D convolution ops.
Definition: Transforms.h:1518
LogicalResult matchAndRewrite(Conv2DOp convOp, PatternRewriter &rewriter) const override
Definition: Transforms.h:1524
FailureOr< Conv1DOp > returningMatchAndRewrite(Conv2DOp convOp, PatternRewriter &rewriter) const
SmallVector< Value > replacements
Definition: Transforms.h:542
Fuse two linalg.generic operations that have a producer-consumer relationship captured through fusedO...
Definition: Transforms.h:551
llvm::DenseMap< Value, Value > replacements
Definition: Transforms.h:553
Rewrite extract_slice(tensor.pad(x)) into tensor.pad(extract_slice(x)).
Definition: Transforms.h:1765
std::function< std::optional< bool >(tensor::ExtractSliceOp)> ControlFn
A function to control pattern application and rewrite logic.
Definition: Transforms.h:1775
LogicalResult matchAndRewrite(tensor::ExtractSliceOp sliceOp, PatternRewriter &rewriter) const override
Definition: Transforms.cpp:992
ExtractSliceOfPadTensorSwapPattern(MLIRContext *context, ControlFn controlFn=nullptr, PatternBenefit benefit=1)
Definition: Transforms.h:1777
Transformation information returned after reduction tiling.
Definition: Transforms.h:992
SmallVector< Operation * > mergeOps
The final reduction operation merging all the partial reductions.
Definition: Transforms.h:996
SmallVector< Value > initialValues
Initial values used for partial reductions.
Definition: Transforms.h:998
scf::ForallOp loops
The scf.forall operation that iterate over the tiles.
Definition: Transforms.h:1000
SmallVector< Operation * > parallelTiledOps
The partial reduction tiled op generated.
Definition: Transforms.h:994
Match and rewrite for the pattern:
Definition: Transforms.h:1728
LogicalResult matchAndRewrite(vector::TransferReadOp xferOp, PatternRewriter &rewriter) const override
TODO: use interfaces, side-effects and aliasing analysis as appropriate, when available.
Match and rewrite for the pattern:
Definition: Transforms.h:1756
LogicalResult matchAndRewrite(vector::TransferWriteOp xferOp, PatternRewriter &rewriter) const override
TODO: use interfaces, side-effects and aliasing analysis as appropriate, when available.
Linalg generalization pattern.
Definition: Transforms.h:1575
LogicalResult matchAndRewrite(LinalgOp op, PatternRewriter &rewriter) const override
Definition: Transforms.h:1585
FailureOr< GenericOp > returningMatchAndRewrite(LinalgOp op, PatternRewriter &rewriter) const
matchAndRewrite implementation that returns the significant transformed pieces of IR.
Definition: Transforms.h:1581
Options that allow distribution of loops generated in Linalg transforms to processors while generatin...
Definition: Utils.h:319
SmallVector< Attribute > paddingValues
A padding value for every operand.
Definition: Transforms.h:283
LinalgPaddingOptions & setPadToMultipleOf(ArrayRef< int64_t > m)
Definition: Transforms.h:296
DenseMap< std::pair< unsigned, unsigned >, OpFoldResult > sizeToPadTo
A mapping between an operand and shape dim, and a size for a padding dimension.
Definition: Transforms.h:304
std::optional< SmallVector< int64_t > > padToMultipleOf
A list of multiples to which each padding dimension should be padded to.
Definition: Transforms.h:295
OpFoldResult getSizeToPadTo(unsigned operandIndex, unsigned dimIndex) const
Given the operand index and shape dim it returns the size to pad to.
Definition: Transforms.h:312
LinalgPaddingOptions & setNofoldFlags(ArrayRef< bool > pp)
Definition: Transforms.h:320
LinalgPaddingOptions & setPaddingDimensions(ArrayRef< int64_t > pd)
Definition: Transforms.h:290
LinalgPaddingOptions & setTransposePaddings(ArrayRef< SmallVector< int64_t >> tp)
Definition: Transforms.h:334
SmallVector< SmallVector< int64_t > > transposePaddings
A permutation vector for every operand used to transpose the packed PadOp results.
Definition: Transforms.h:332
LinalgPaddingOptions & setSizeToPadTo(unsigned operandIndex, unsigned dimIndex, OpFoldResult size)
Definition: Transforms.h:305
LinalgPaddingOptions & setPaddingValues(ArrayRef< Attribute > pv)
Definition: Transforms.h:284
SmallVector< bool > nofoldFlags
A flag for every operand to mark the PadOp as nofold which enables packing for statically shaped oper...
Definition: Transforms.h:319
LinalgPaddingOptions & setCopyBackOp(CopyBackOp op)
Definition: Transforms.h:346
LinalgPaddingOptions & setHoistPaddings(ArrayRef< int64_t > hp)
Definition: Transforms.h:326
SmallVector< int64_t > hoistPaddings
A number of loops to hoist the PadOp out for every operand.
Definition: Transforms.h:325
SmallVector< int64_t > paddingDimensions
A list of iterator dimensions to pad.
Definition: Transforms.h:289
CopyBackOp copyBackOp
The op to be used for copying the padded result to the original destination tensor.
Definition: Transforms.h:345
std::optional< unsigned > alignment
Alignment of promoted buffer. If std::nullopt do not specify alignment.
Definition: Transforms.h:433
LinalgPromotionOptions & setUseFullTileBuffersByDefault(bool use)
Definition: Transforms.h:421
bool useAlloca
Use alloca with the default allocation scheme.
Definition: Transforms.h:446
LinalgPromotionOptions & setAlignment(unsigned align)
Definition: Transforms.h:434
std::optional< Attribute > memorySpace
Memory space of promoted buffer.
Definition: Transforms.h:440
bool useOriginalSubviewSize
If true, buffers will be allocated with the original subview size.
Definition: Transforms.h:427
std::optional< CopyCallbackFn > copyOutFn
Definition: Transforms.h:466
std::optional< CopyCallbackFn > copyInFn
Callback function to do the copy of data to and from the promoted subview.
Definition: Transforms.h:465
LinalgPromotionOptions & setUseAlloca(bool use)
Definition: Transforms.h:447
std::optional< DenseSet< unsigned > > operandsToPromote
Indices of subViews to promote.
Definition: Transforms.h:398
LinalgPromotionOptions & setCopyInOutFns(CopyCallbackFn const &copyIn, CopyCallbackFn const &copyOut)
Definition: Transforms.h:467
LinalgPromotionOptions & setUseFullTileBuffers(ArrayRef< bool > useFullTiles)
Definition: Transforms.h:410
std::optional< AllocBufferCallbackFn > allocationFn
Callback function to do the allocation of the promoted buffer.
Definition: Transforms.h:454
bool useFullTileBuffersDefault
If true all operands unspecified by useFullTileBuffers will use the full view, otherwise the partial ...
Definition: Transforms.h:420
std::optional< DeallocBufferCallbackFn > deallocationFn
Definition: Transforms.h:455
LinalgPromotionOptions & setMemorySpace(Attribute memorySpc)
Definition: Transforms.h:441
LinalgPromotionOptions & setAllocationDeallocationFns(AllocBufferCallbackFn const &allocFn, DeallocBufferCallbackFn const &deallocFn)
Definition: Transforms.h:457
LinalgPromotionOptions & setUseOriginalSubviewSize(bool originalSize)
Definition: Transforms.h:428
std::optional< llvm::SmallBitVector > useFullTileBuffers
If ith element of useFullTiles is true the full view should be used for the promoted buffer of the it...
Definition: Transforms.h:409
LinalgPromotionOptions & setOperandsToPromote(ArrayRef< int64_t > operands)
Definition: Transforms.h:399
LogicalResult matchAndRewrite(GenericOp op, PatternRewriter &rewriter) const override
Definition: Transforms.h:1599
FailureOr< GenericOp > returningMatchAndRewrite(GenericOp op, PatternRewriter &rewriter) const
Definition: Transforms.h:1595
std::optional< LinalgLoopDistributionOptions > tileDistribution
When specified, specifies distribution of generated tile loops to processors.
Definition: Transforms.h:273
LinalgTilingAndFusionOptions & setTileSizes(ArrayRef< int64_t > ts)
Definition: Transforms.h:265
SmallVector< int64_t > tileInterchange
Tile interchange used to permute the tile loops.
Definition: Transforms.h:270
LinalgTilingAndFusionOptions & setDistributionOptions(LinalgLoopDistributionOptions distributionOptions)
Definition: Transforms.h:275
SmallVector< int64_t > tileSizes
Tile sizes used to tile the root operation.
Definition: Transforms.h:264
LinalgTilingOptions & setLoopType(LinalgTilingLoopType lt)
Definition: Transforms.h:229
LinalgTilingOptions & setDistributionTypes(ArrayRef< StringRef > types)
Definition: Transforms.h:247
LinalgTilingOptions & setInterchange(ArrayRef< unsigned > interchange)
Definition: Transforms.h:221
LinalgTilingLoopType loopType
The type of tile loops to generate.
Definition: Transforms.h:227
LinalgTilingOptions & setTileSizeComputationFunction(TileSizeComputationFunction fun)
Definition: Transforms.h:198
LinalgTilingOptions & setTileSizes(const SmallVector< Value, 4 > &ts)
Set the tileSizeComputationFunction to return the values ts.
Definition: Transforms.h:205
LinalgTilingOptions & setPeeledLoops(ArrayRef< int64_t > loops)
Definition: Transforms.h:255
SmallVector< int64_t > peeledLoops
Peel the specified loops.
Definition: Transforms.h:253
LinalgTilingOptions & setDistributionOptions(LinalgLoopDistributionOptions distributionOptions)
Definition: Transforms.h:239
SmallVector< unsigned, 4 > interchangeVector
The interchange vector to reorder the tiled loops.
Definition: Transforms.h:219
TileSizeComputationFunction tileSizeComputationFunction
Computation function that returns the tile sizes for each operation.
Definition: Transforms.h:195
LinalgTilingOptions & scalarizeDynamicDims()
Tile all dynamic dimensions by 1.
std::optional< LinalgLoopDistributionOptions > distribution
When specified, specifies distribution of generated tile loops to processors.
Definition: Transforms.h:236
SmallVector< StringRef, 2 > distributionTypes
Specification markers of how to distribute the linalg.tiled_loop.
Definition: Transforms.h:245
linalg::TransposeOp transposeOp
Definition: Transforms.h:1221
tensor::ExpandShapeOp expandShapeOp
Definition: Transforms.h:1220
tensor::ExtractSliceOp extractSliceOp
Definition: Transforms.h:1233
linalg::TransposeOp transposeOp
Definition: Transforms.h:1231
tensor::CollapseShapeOp collapseShapeOp
Definition: Transforms.h:1232
A description of a multi-size tiling comprising tile sizes and numbers of tiles, expressed as Values ...
Definition: Transforms.h:939
Struct to hold the result of a pack call.
Definition: Transforms.h:1242
SmallVector< linalg::UnPackOp > unPackOps
Definition: Transforms.h:1245
linalg::LinalgOp packedLinalgOp
Definition: Transforms.h:1244
SmallVector< linalg::PackOp > packOps
Definition: Transforms.h:1243
Struct to hold the result of a packTranspose call.
Definition: Transforms.h:1254
linalg::LinalgOp transposedLinalgOp
Definition: Transforms.h:1256
linalg::UnPackOp transposedUnPackOp
Definition: Transforms.h:1257
PadTilingInterfaceOptions & setPaddingSizes(ArrayRef< OpFoldResult > m)
Definition: Transforms.h:361
SmallVector< Attribute > paddingValues
A padding value for every operand.
Definition: Transforms.h:354
PadTilingInterfaceOptions & setPadToMultipleOf(bool b)
Definition: Transforms.h:368
bool padToMultipleOf
Pad iterator paddingDimension[i] to next multiple of paddingSizes[i] if true.
Definition: Transforms.h:367
PadTilingInterfaceOptions & setPaddingValues(ArrayRef< Attribute > pv)
Definition: Transforms.h:355
SmallVector< OpFoldResult > paddingSizes
A list of iterator dimensions sizes to pad to.
Definition: Transforms.h:360
Create a new buffer using the allocationFn provided.
Definition: Transforms.h:811
Split Reduction options.
Definition: Transforms.h:476
Apply transformation to split the single linalg op reduction into a parallel and reduction dimension.
Definition: Transforms.h:1130
Perform standalone tiling of a single LinalgOp by tileSizes.
Definition: Transforms.h:771
SmallVector< Operation *, 8 > loops
Definition: Transforms.h:773
SmallVector< Value, 4 > tensorResults
Definition: Transforms.h:774
Transformation information returned after vectorizing.
Definition: Transforms.h:868
SmallVector< Value > replacements
Results of the vectorization transform to replace the original operation.
Definition: Transforms.h:870
SmallVector< T > tripCounts
Number of tiles associated with each size.
Definition: Transforms.h:930
T lowTripCount
Number of tiles associated with each size.
Definition: Transforms.h:922
Helper struct to hold the results of building a packing loop nest.
Definition: Transforms.h:641
SmallVector< OpFoldResult > strides
Definition: Transforms.h:642
SmallVector< Value > leadingPackedTensorIndexings
Definition: Transforms.h:643
SmallVector< Value > clonedLoopIvs
Definition: Transforms.h:643
SmallVector< OpFoldResult > sizes
Definition: Transforms.h:642
SmallVector< OpFoldResult > offsets
Definition: Transforms.h:642