MLIR  19.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/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 } // namespace bufferization
35 
36 namespace linalg {
37 
38 class LinalgOp;
39 
40 //===----------------------------------------------------------------------===//
41 // Utils.
42 //===----------------------------------------------------------------------===//
43 
44 /// Return vector::CombiningKind for the given op.
45 std::optional<vector::CombiningKind> getCombinerOpKind(Operation *combinerOp);
46 
47 //===----------------------------------------------------------------------===//
48 // Bufferization-related transforms.
49 //===----------------------------------------------------------------------===//
50 
52  enum class AllocOp { MemrefAlloc = 0, MemrefAlloca = 1 };
54 
55  enum class MemcpyOp {
57  MemrefCopy = 1,
58  LinalgCopy = 2
59  };
61 
62  /// If set to "true", only the destination tensor operands are bufferized to
63  /// a new allocation (and wrapped in "bufferization.to_tensor"), but not the
64  /// targeted op itself.
66 
67  /// If set to "true", a `memref.dealloc` operation will be emitted for each
68  /// allocated buffer. Otherwise, the memory is leaked, which is useful if
69  /// the buffer deallocation pipeline should be run after bufferization is
70  /// done.
71  bool emitDealloc = false;
72 };
73 
74 /// Materialize a buffer allocation for the given tensor.pad op and lower the
75 /// op to linalg.fill/linalg.generic + bufferization.materialize_in_destination.
76 /// E.g.:
77 ///
78 /// %0 = tensor.pad low[%l] high[%h] %t ...
79 ///
80 /// is lowered to:
81 ///
82 /// %alloc = memref.alloc
83 /// linalg.fill ... outs(%alloc)
84 /// %subview = memref.subview %alloc [%l] [...] [1]
85 /// bufferization.materialize_in_destination %t in %subview
86 /// %0 = bufferization.to_tensor %alloc restrict writable
87 ///
88 /// In addition to rewriting the IR as shown above, this function returns the
89 /// newly allocated buffer. The `insertionPoint` parameter can be used to
90 /// specify a custom insertion point for the buffer allocation.
93  tensor::PadOp padOp, Attribute memorySpace = {},
94  Operation *insertionPoint = nullptr);
95 
96 /// Materialize a buffer allocation for the given vector.mask op and bufferize
97 /// the op, including its region. E.g.:
98 ///
99 /// %0 = vector.mask {
100 /// vector.transfer_write %v, %t : vector<16xf32>, tensor<?xf32>
101 /// } : vector<16xi1> -> tensor<?xf32>
102 ///
103 /// is lowered to:
104 ///
105 /// %alloc = memref.alloc
106 /// bufferization.materialize_in_destination %t in %subview
107 /// vector.mask {
108 /// vector.transfer_write %arg0, %alloc : vector<16xf32>, memref<?xf32>
109 /// } : vector<16xi1>
110 /// %0 = bufferization.to_tensor %alloc restrict writable
111 ///
112 /// In addition to rewriting the IR as shown above, this function returns the
113 /// newly allocated buffer. The `insertionPoint` parameter can be used to
114 /// specify a custom insertion point for the buffer allocation.
116  const BufferizeToAllocationOptions &options,
117  vector::MaskOp maskOp, Attribute memorySpace = {},
118  Operation *insertionPoint = nullptr);
119 
120 /// Materialize a buffer allocation for the given bufferization.alloc_tensor op
121 /// and lower the op to memref.alloc + memref.tensor_store.
122 ///
123 /// In addition to rewriting the IR, this function returns the newly allocated
124 /// buffer. The `insertionPoint` parameter can be used to specify a custom
125 /// insertion point for the buffer allocation.
126 Value bufferizeToAllocation(RewriterBase &rewriter,
127  const BufferizeToAllocationOptions &options,
128  bufferization::AllocTensorOp allocTensorOp,
129  Attribute memorySpace = {},
130  Operation *insertionPoint = nullptr);
131 
132 /// Bufferize the given op with tensor semantics and materialize the result in
133 /// a newly allocated buffer.
134 ///
135 /// Only bufferizable ops that bufferize to a memory write or have an
136 /// aliasing OpOperand (and do not themselves bufferize to an allocation) are
137 /// supported. They are bufferized using their BufferizableOpInterface
138 /// implementation.
139 ///
140 /// Selected ops that bufferize to an allocation (or need special handling) are
141 /// also supported:
142 /// - tensor.pad
143 /// - vector.mask
144 ///
145 /// This function returns the newly allocated buffer. The `insertionPoint`
146 /// parameter can be used to specify a custom insertion point for the buffer
147 /// allocation.
148 Value bufferizeToAllocation(RewriterBase &rewriter,
149  const BufferizeToAllocationOptions &options,
150  Operation *op, Attribute memorySpace = {},
151  Operation *insertionPoint = nullptr);
152 
153 /// Try to eliminate tensor::EmptyOps inside `op` that are anchored on a
154 /// LinalgOp. This transforms looks for LinalgOps that have an unused output
155 /// operand and an input operand that is rooted in a tensor::EmptyOp. The
156 /// tensor::EmptyOp uses are replaced with the output operand and the two
157 /// operands of the LinalgOp are swapped.
158 ///
159 /// Example:
160 /// %0 = tensor.empty()
161 /// %1 = linalg.matmul ins(...) outs(%0)
162 /// %2 = linalg.generic ins(%1) outs(%dest) {
163 /// ^bb0(%in: f32, %out: f32):
164 /// // out not used
165 /// }
166 ///
167 /// The IR is transformed as follows:
168 /// %0 = tensor.empty()
169 /// %1 = linalg.matmul ins(...) outs(%dest)
170 /// %2 = linalg.generic ins(%0) outs(%1) {
171 /// ^bb0(%in: f32, %out: f32):
172 /// // Use %out instead of %in
173 /// }
174 ///
175 /// The "ins" operand has no uses inside the body of the LinalgOp and can be
176 /// folded away with existing cleanup patterns. Afterwards, the tensor::EmptyOp
177 /// can also fold away.
179  RewriterBase &rewriter, Operation *op,
180  bufferization::OneShotAnalysisState &state);
181 
182 //===----------------------------------------------------------------------===//
183 // Structs that configure the behavior of various transformations.
184 //===----------------------------------------------------------------------===//
185 
187  std::function<SmallVector<Value, 4>(OpBuilder &, Operation *)>;
188 
190  /// Computation function that returns the tile sizes for each operation.
191  /// Delayed construction of constant tile sizes should occur to interoperate
192  /// with folding.
194 
197  tileSizeComputationFunction = std::move(fun);
198  return *this;
199  }
200  /// Set the `tileSizeComputationFunction` to return the values `ts`. The
201  /// values must not fold away when tiling. Otherwise, use a more robust
202  /// `tileSizeComputationFunction`.
204  tileSizeComputationFunction = [=](OpBuilder &, Operation *) { return ts; };
205  return *this;
206  }
207  /// Convenience function to set the `tileSizeComputationFunction` to a
208  /// function that computes tile sizes at the point they are needed. Allows
209  /// proper interaction with folding.
211 
212  /// Tile all dynamic dimensions by 1. I.e., scalarize those dimensions.
213  /// Note: `scalarizeDynamicDims` and `setTileSizes` cannot be used together.
215 
216  /// The interchange vector to reorder the tiled loops.
218 
220  interchangeVector.assign(interchange.begin(), interchange.end());
221  return *this;
222  }
223 
224  /// The type of tile loops to generate.
226 
228  loopType = lt;
229  return *this;
230  }
231 
232  /// When specified, specifies distribution of generated tile loops to
233  /// processors.
234  std::optional<LinalgLoopDistributionOptions> distribution;
235 
238  distribution = std::move(distributionOptions);
239  return *this;
240  }
241 
242  /// Specification markers of how to distribute the `linalg.tiled_loop`.
244 
246  distributionTypes.assign(types.begin(), types.end());
247  return *this;
248  }
249 
250  /// Peel the specified loops.
252 
254  peeledLoops.clear();
255  peeledLoops.append(loops.begin(), loops.end());
256  return *this;
257  }
258 };
259 
261  /// Tile sizes used to tile the root operation.
264  tileSizes.assign(ts.begin(), ts.end());
265  return *this;
266  }
267  /// Tile interchange used to permute the tile loops.
269  /// When specified, specifies distribution of generated tile loops to
270  /// processors.
271  std::optional<LinalgLoopDistributionOptions> tileDistribution;
274  tileDistribution = std::move(distributionOptions);
275  return *this;
276  }
277 };
278 
280  /// A padding value for every operand.
283  paddingValues.assign(pv.begin(), pv.end());
284  return *this;
285  }
286  /// A list of iterator dimensions to pad.
289  paddingDimensions.assign(pd.begin(), pd.end());
290  return *this;
291  }
292  /// A list of multiples to which each padding dimension should be padded to.
293  std::optional<SmallVector<int64_t>> padToMultipleOf;
295  padToMultipleOf.emplace(m.begin(), m.end());
296  return *this;
297  }
298  /// A flag for every operand to mark the PadOp as nofold which enables
299  /// packing for statically shaped operands.
302  packPaddings.assign(pp.begin(), pp.end());
303  return *this;
304  }
305  /// A number of loops to hoist the PadOp out for every operand.
308  hoistPaddings.assign(hp.begin(), hp.end());
309  return *this;
310  }
311  /// A permutation vector for every operand used to transpose the packed
312  /// PadOp results.
316  transposePaddings.assign(tp.begin(), tp.end());
317  return *this;
318  }
319  enum class CopyBackOp : int8_t {
320  None = 0,
322  LinalgCopy = 2
323  };
324  /// The op to be used for copying the padded result to the original
325  /// destination tensor.
328  copyBackOp = op;
329  return *this;
330  }
331 };
332 
333 /// Callback function type used to perform the allocation for the promoted
334 /// `subView`. In `boundingSubViewsize` a best attempt is made to find the
335 /// smallest constant value for the size of the buffer needed for each
336 /// dimension. If that is not possible, contains the dynamic size of the
337 /// subview. The call back should return the buffer to use.
338 using AllocBufferCallbackFn = std::function<std::optional<Value>(
339  OpBuilder &b, memref::SubViewOp subView,
340  ArrayRef<Value> boundingSubViewSize, DataLayout &layout)>;
341 
342 /// Callback function type used to deallocate the buffers used to hold the
343 /// promoted subview.
345  std::function<LogicalResult(OpBuilder &b, Value buffer)>;
346 
347 /// Callback function type used to insert copy from original subview to
348 /// subview of the promoted region for the read operands/subview of promoted
349 /// region to original subview for the results. The copy has to happen from
350 /// `src` to `dst`.
352  std::function<LogicalResult(OpBuilder &b, Value src, Value dst)>;
353 
355  /// Indices of subViews to promote. If `std::nullopt`, try to promote all
356  /// operands.
357  std::optional<DenseSet<unsigned>> operandsToPromote;
360  operandsToPromote->insert(operands.begin(), operands.end());
361  return *this;
362  }
363  /// If ith element of `useFullTiles` is true the full view should be used
364  /// for the promoted buffer of the ith operand in `operandsToPromote`.
365  /// Otherwise the partial view will be used. The decision is defaulted to
366  /// `useFullTileBuffersDefault` when `useFullTileBuffers` is std::nullopt and
367  /// for operands missing from `useFullTileBuffers`.
368  std::optional<llvm::SmallBitVector> useFullTileBuffers;
370  unsigned size = useFullTiles.size();
371  llvm::SmallBitVector tmp(size, false);
372  for (unsigned i = 0; i < size; ++i)
373  tmp[i] = useFullTiles[i];
374  useFullTileBuffers = tmp;
375  return *this;
376  }
377  /// If true all operands unspecified by `useFullTileBuffers` will use the
378  /// full view, otherwise the partial view.
382  return *this;
383  }
384  /// Alignment of promoted buffer. If `std::nullopt` do not specify alignment.
385  std::optional<unsigned> alignment;
387  alignment = align;
388  return *this;
389  }
390  /// Memory space of promoted buffer. If `std::nullopt` do not specify memory
391  /// space.
392  std::optional<Attribute> memorySpace;
394  memorySpace = memorySpc;
395  return *this;
396  }
397  /// Use alloca with the default allocation scheme.
398  bool useAlloca = false;
400  useAlloca = use;
401  return *this;
402  }
403  /// Callback function to do the allocation of the promoted buffer. If
404  /// std::nullopt, then the default allocation scheme of allocating a
405  /// memref<?xi8> buffer followed by a view operation is used.
406  std::optional<AllocBufferCallbackFn> allocationFn;
407  std::optional<DeallocBufferCallbackFn> deallocationFn;
410  DeallocBufferCallbackFn const &deallocFn) {
411  allocationFn = allocFn;
412  deallocationFn = deallocFn;
413  return *this;
414  }
415  /// Callback function to do the copy of data to and from the promoted
416  /// subview. If std::nullopt then a memref.copy is used.
417  std::optional<CopyCallbackFn> copyInFn;
418  std::optional<CopyCallbackFn> copyOutFn;
420  CopyCallbackFn const &copyOut) {
421  copyInFn = copyIn;
422  copyOutFn = copyOut;
423  return *this;
424  }
425 };
426 
427 /// Split Reduction options.
429  // Ratio used to split the reduction dimension. If the ratio is <= 1,
430  // nothing will be done.
431  int64_t ratio = 0;
432  // Index where the extra dimension is added to the intermediate tensor
433  // shape.
434  unsigned index = 0;
435  // If the inner dimension after splitting is parallel or reduction.
436  bool innerParallel = false;
437 };
438 
439 /// Function signature to control reduction splitting. This returns
440 /// `SplitReductionOptions`.
441 // TODO: don't use unsigned unless doing bit manipulation.
443  std::function<SplitReductionOptions(LinalgOp op)>;
444 
445 //===----------------------------------------------------------------------===//
446 // Preconditions that ensure the corresponding transformation succeeds and can
447 // be applied as a rewrite pattern.
448 //===----------------------------------------------------------------------===//
449 
450 /// Return true if two `linalg.generic` operations with producer/consumer
451 /// relationship through `fusedOperand` can be fused using elementwise op
452 /// fusion.
453 bool areElementwiseOpsFusable(OpOperand *fusedOperand);
454 
455 /// Promote memref.subviews feeding linalg-on-buffers operations.
458 
459 /// Return success if the operation can be vectorized.
461  ArrayRef<int64_t> inputVectorSizes = {},
462  ArrayRef<bool> inputScalableVecDims = {},
463  bool vectorizeNDExtract = false,
464  bool flatten1DDepthwiseConv = false);
465 
466 //===----------------------------------------------------------------------===//
467 // Transformations exposed as functional-style API calls.
468 //===----------------------------------------------------------------------===//
469 
471 
472 /// Transformation to drop unit-extent dimensions from `linalg.generic`
473 /// operations.
476 
479 
480  using ControlFnTy = std::function<SmallVector<unsigned>(Operation *)>;
482  if (auto genericOp = dyn_cast_or_null<GenericOp>(op)) {
483  return llvm::to_vector(llvm::seq<unsigned>(0, genericOp.getNumLoops()));
484  }
485  if (auto padOp = dyn_cast_or_null<tensor::PadOp>(op)) {
486  return llvm::to_vector(
487  llvm::seq<unsigned>(0, padOp.getSourceType().getRank()));
488  }
489  return SmallVector<unsigned>{};
490  };
491 };
492 LogicalResult dropUnitDims(RewriterBase &rewriter, GenericOp genericOp,
493  const ControlDropUnitDims &options);
494 
495 /// Fuse two `linalg.generic` operations that have a producer-consumer
496 /// relationship captured through `fusedOperand`. The method expects
497 /// that `areElementwiseOpsFusable` returns true for the given `fusedOperand`.
501  static llvm::SmallDenseSet<int>
502  getPreservedProducerResults(GenericOp producer, GenericOp consumer);
503 };
505 fuseElementwiseOps(RewriterBase &rewriter, OpOperand *fusedOperand);
506 
507 /// Try to peel and canonicalize loop `op` and return the new result.
508 /// Also applies affine_min/max bounds simplification on the fly where relevant.
509 // TODO: Add support for scf.parallel and affine.for loops.
511 
512 /// Peel 'loops' and applies affine_min/max bounds simplification on the fly
513 /// where relevant.
514 void peelLoops(RewriterBase &rewriter, ArrayRef<scf::ForOp> loops);
515 
516 /// Pad the iterator dimensions `paddingDimensions` of all `opToPad` operands
517 /// to a static bounding box. The original `opToPad` is cloned and operates on
518 /// the padded tensors.
519 ///
520 /// * "options.padToMultipleOf" indicates that each padding dimension should be
521 /// padded to the specified multiple.
522 /// * Use "options.paddingValues" and "options.packPaddings" to set padding
523 /// value and nofold attribute of the created tensor::PadOps, respectively.
524 /// * The unpadded results (extracted slice of the cloned operation) are
525 /// returned via `replacements`.
526 /// * The tensor::PadOps are returned via `padOps`.
527 /// * "options.copyBackOp" specifies the op type for copying back the unpadded
528 /// result to the original destination tensor.
529 LogicalResult rewriteAsPaddedOp(RewriterBase &rewriter, LinalgOp opToPad,
531  LinalgOp &paddedOp,
532  SmallVector<Value> &replacements,
534 
535 namespace detail {
536 
537 /// Helper struct to hold the results of building a packing loop nest.
541  GenericOp maybeTransposeOp;
542  tensor::PadOp hoistedPadOp;
543 };
544 
545 /// Build the packing loop nest required to hoist `opToHoist` above
546 /// `outermostEnclosingForOp`.
547 /// The loop nest is built just before `outermostEnclosingForOp`.
549 buildPackingLoopNest(RewriterBase &rewriter, tensor::PadOp opToHoist,
550  scf::ForOp outermostEnclosingForOp,
551  ArrayRef<int64_t> transposeVector);
552 
553 } // namespace detail
554 
555 /// Mechanically hoist padding operations on tensors by `numLoops` into a new,
556 /// generally larger tensor. This achieves packing of multiple padding ops into
557 /// a larger tensor. On success, `opToHoist` is replaced by the cloned version
558 /// in the packing loop so the caller can continue reasoning about the padding
559 /// operation. If `transposeVector` is non-empty, hoist padding introduces a
560 /// GenericOp to transpose the padded tensor before inserting it into the packed
561 /// tensor. A `transposeVector` can change the storage order of the padded
562 /// tensor but does not change the order of the pack or compute loops.
563 ///
564 /// TODO: In the future, we should consider rewriting as a tensor.pack after
565 /// hoisting since this abstraction is now available.
566 ///
567 /// Example in pseudo-mlir:
568 /// =======================
569 ///
570 /// If hoistPaddingOnTensors is called with `nLoops` = 2 on the following IR.
571 /// ```
572 /// scf.for (%i, %j, %k)
573 /// %st0 = tensor.extract_slice f(%i, %k) : ... to tensor<?x?xf32>
574 /// %0 = tensor.pad %st0 low[0, 0] high[...] {
575 /// ^bb0( ... ):
576 /// linalg.yield %pad
577 /// } : tensor<?x?xf32> to tensor<4x8xf32>
578 /// compute(%0)
579 /// ```
580 ///
581 /// IR resembling the following is produced:
582 ///
583 /// ```
584 /// scf.for (%i) {
585 /// %packed_init = tensor.empty range(%j) : tensor<?x4x8xf32>
586 /// %packed = scf.for (%k) iter_args(%p : %packed_init) {
587 /// %st0 = tensor.extract_slice f(%i, %k) : ... to tensor<?x?xf32>
588 /// %0 = tensor.pad %st0 low[0, 0] high[...] {
589 /// ^bb0( ... ):
590 /// linalg.yield %pad
591 /// } : tensor<?x?xf32> to tensor<4x8xf32>
592 /// %1 = tensor.insert_slice %0 ...
593 /// : tensor<4x8xf32> to tensor<?x4x8xf32>
594 /// scf.yield %1: tensor<?x4x8xf32>
595 /// } -> tensor<?x4x8xf32>
596 /// scf.for (%j, %k) {
597 /// %st0 = tensor.extract_slice %packed [%k, 0, 0][1, 4, 8][1, 1, 1] :
598 /// tensor<?x4x8xf32> to tensor<4x8xf32>
599 /// compute(%st0)
600 /// }
601 /// }
602 /// ```
604 hoistPaddingOnTensors(RewriterBase &rewriter, tensor::PadOp opToHoist,
605  int64_t numLoops, ArrayRef<int64_t> transposeVector,
606  tensor::PadOp &hoistedOp,
607  SmallVectorImpl<GenericOp> &transposeOps);
608 /// Calls into `hoistPaddingOnTensors` with a local IRRewriter.
610 hoistPaddingOnTensors(tensor::PadOp opToHoist, int64_t numLoops,
611  ArrayRef<int64_t> transposeVector,
612  tensor::PadOp &hoistedOp,
613  SmallVectorImpl<GenericOp> &transposeOps);
614 
615 /// Apply padding and hoisting to `linalgOp` according to the configuration
616 /// specified in `options`.
618  LinalgOp linalgOp,
620 
621 /// Split the given `op` into two parts along the given iteration space
622 /// `dimension` at the specified `splitPoint`, and return the two parts.
623 /// If the second part is statically known to be empty, do not create it
624 /// and return nullptr instead. Error state is signalled by returning
625 /// a pair of nullptrs.
626 ///
627 /// For example, the following op:
628 ///
629 /// linalg.matmul ins(%0, %1 : tensor<128x32xf32>, tensor<32x64xf32>)
630 /// outs(%2 : tensor<128x64xf32>)
631 ///
632 /// split along the first dimension at position 42 will result in:
633 ///
634 /// %3 = tensor.extract_slice %0[0, 0][42, 32][1, 1]
635 /// %4 = tensor.extract_slice %2[0, 0][42, 64][1, 1]
636 /// %5 = linalg.matmul ins(%3, %1 : tensor<42x32xf32>, tensor<32x64xf32>)
637 /// outs(%5 : tensor<42x64xf32>)
638 /// %6 = tensor.insert_slice %5 into %2[0, 0][42, 64][1, 1]
639 ///
640 /// %7 = tensor.extract_slice %0[42, 0][86, 32][1, 1]
641 /// %8 = tensor.extract_slice %6[42, 0][86, 64][1, 1]
642 /// %9 = linalg.matmul ins(%7, %1 : tensor<86x32xf32>, tensor<32x64xf32>)
643 /// outs(%8 : tensor<86x64xf32>)
644 /// tensor.insert_slice %5 into %6[42, 0][86, 64][1, 1]
645 ///
646 /// Note that there is no simplification other than constant propagation applied
647 /// to slice extraction and insertion.
648 std::pair<TilingInterface, TilingInterface> splitOp(RewriterBase &rewriter,
649  TilingInterface op,
650  unsigned dimension,
651  OpFoldResult splitPoint);
652 
653 /// Perform standalone tiling of a single LinalgOp by `tileSizes`.
654 /// and permute the loop nest according to `interchangeVector`
655 /// The permutation is expressed as a list of integers that specify
656 /// the new ordering of the loop nest. The length of `interchangeVector`
657 /// must be equal to the length of `tileSizes`.
658 /// An empty vector is interpreted as the identity permutation and the
659 /// transformation returns early.
660 ///
661 /// Return a struct containing the tiled loops in the specified order
662 /// and the cloned op if successful, std::nullopt otherwise.
663 ///
664 /// E.g. the permutation `(i,j,k) -> (j,k,i)` is expressed by
665 /// `interchangeVector = [1,2,0]`. All values in `interchangeVector` must be
666 /// integers, in the range 0..`tileSizes.size()` without duplications
667 /// (i.e. `[1,1,2]` is an invalid permutation).
669  LinalgOp op;
672 };
675 
676 /// Interchange the `iterator_types` and `iterator_maps` dimensions and adapts
677 /// the index accesses of `op`. This is an in-place transformation controlled
678 /// by `interchangeVector`. An empty vector is interpreted as the identity
679 /// permutation and the transformation returns early.
680 ///
681 /// E.g. the permutation `(i,j,k) -> (j,k,i)` is expressed with
682 /// `interchangeVector = [1,2,0]`. All values in `interchangeVector` must be
683 /// integers, in the range 0..`op.rank` without duplications
684 /// (i.e. `[1,1,2]` is an invalid permutation).
685 ///
686 /// Return failure if the permutation is not valid.
688  GenericOp genericOp,
689  ArrayRef<unsigned> interchangeVector);
690 
691 /// Create a GenericOp from the given named operation `namedOp` and replace
692 /// namedOp.
693 /// Return failure if `namedOp` is a GenericOp or misses a region builder.
695  LinalgOp namedOp);
696 
697 /// Create a namedOp from the given GenericOp and replace the GenericOp.
698 /// Currently we can specialize only trivial linalg copy operations.
700  GenericOp genericOp);
701 
702 /// Create a new buffer using the `allocationFn` provided. The size of this
703 /// buffer is the smallest constant bounding size along each dimension that
704 /// can be computed for the size of the result of `subView`. Returns the
705 /// allocated buffer as `fullLocalView` and the view that matches the size of
706 /// the result of subview operation as `partialLocalView`.
710 };
712 promoteSubviewAsNewBuffer(OpBuilder &b, Location loc, memref::SubViewOp subView,
713  const AllocBufferCallbackFn &allocationFn,
714  DataLayout &layout);
715 
716 /// Promote the `subViews` into a new buffer allocated at the insertion point
717 /// `b`. Promotion occurs in 3 steps:
718 /// 1. Create a new buffer for a full tile (i.e. not clipped at the
719 /// boundary).
720 /// 2. Take a full view on the buffer.
721 /// 3. Take a partial slice of the full view in step 2. and copy into it.
722 ///
723 /// Return the modified linalg op (the modification happens in place) as well
724 /// as all the copy ops created.
727 
728 /// Allocate the subview in the GPU workgroup memory.
729 std::optional<Value> allocateWorkgroupMemory(OpBuilder &builder,
730  memref::SubViewOp subview,
731  ArrayRef<Value> sizeBounds,
732  DataLayout &);
733 
734 /// In case of GPU group memory there is no need to deallocate.
736 
737 /// Create Memref copy operations and add gpu barrier guards before and after
738 /// the copy operation to ensure data integrity.
740 
741 /// Allocate the subview in the GPU private memory.
742 std::optional<Value> allocateGPUPrivateMemory(OpBuilder &builder,
743  memref::SubViewOp subview,
744  ArrayRef<Value> sizeBounds,
745  DataLayout &);
746 
747 /// Normal copy to between src and dst.
749 
750 /// In case of GPU private memory there is no need to deallocate since the
751 /// memory is freed when going outside of the scope.
753 
754 /// Emit a suitable vector form for an operation. If provided,
755 /// `inputVectorSizes` are used to vectorize this operation. `inputVectorSizes`
756 /// must match the rank of the iteration space of the operation and the sizes
757 /// must be smaller or equal than their counterpart interation space sizes, if
758 /// static. `inputVectorShapes` also allows the vectorization of operations with
759 /// dynamic shapes.
761  ArrayRef<int64_t> inputVectorSizes = {},
762  ArrayRef<bool> inputScalableVecDims = {},
763  bool vectorizeNDExtract = false,
764  bool flatten1DDepthwiseConv = false);
765 
766 /// Emit a suitable vector form for a Copy op with fully static shape.
767 LogicalResult vectorizeCopy(RewriterBase &builder, memref::CopyOp copyOp);
768 
769 /// Emit a loop nest of `scf.for` with the proper body for `linalgOp`.
770 FailureOr<LinalgLoops> linalgOpToLoops(RewriterBase &rewriter,
771  LinalgOp linalgOp);
772 
773 /// Emit a loop nest of `scf.parallel` with the proper body for `linalgOp`.
774 FailureOr<LinalgLoops> linalgOpToParallelLoops(RewriterBase &rewriter,
775  LinalgOp linalgOp);
776 
777 /// Emit a loop nest of `affine.for` with the proper body for `linalgOp`.
778 FailureOr<LinalgLoops> linalgOpToAffineLoops(RewriterBase &rewriter,
779  LinalgOp linalgOp);
780 
781 /// Creates a number of ranges equal to the number of non-zero in `tileSizes`.
782 /// One for each loop of the LinalgOp that is tiled. The `tileSizes` argument
783 /// has one entry per surrounding loop. It uses zero as the convention that a
784 /// particular loop is not tiled. This convention simplifies implementations
785 /// by avoiding affine map manipulations. The returned ranges correspond to
786 /// the loop ranges, in the proper order, that are tiled and for which new
787 /// loops will be created. Also the function returns a map from loop indices
788 /// of the LinalgOp to the corresponding non-empty range indices of newly
789 /// created loops.
791 std::tuple<SmallVector<Range, 4>, LoopIndexToRangeIndexMap>
793  ArrayRef<OpFoldResult> allShapeSizes,
794  ArrayRef<OpFoldResult> allTileSizes);
795 
796 namespace detail {
797 template <typename T>
799  /// Tile sizes.
801  /// Number of tiles associated with each size.
803 };
804 } // namespace detail
805 
806 /// A description of a multi-size tiling comprising tile sizes and numbers of
807 /// tiles, expressed as Values which may or may not be constant. Multi-size
808 /// currently means two-size.
810  : public detail::MultiSizeSpecificationBase<Value> {};
812  : public detail::MultiSizeSpecificationBase<int64_t> {};
813 
814 /// Emits the IR computing the multi-sized tiling specification with two tile
815 /// sizes not exceeding `targetSize`, each divisible by `sizeDivisor`, such
816 /// that there exist numbers of tiles with these sizes that fully cover the
817 /// given iteration space `dimension` of the structured `op`.
818 ///
819 /// The computation is as follows:
820 ///
821 /// b = originalTripCount floordiv sizeDivisor
822 /// t = (targetSize + sizeDivisor - 1) floordiv sizeDivisor
823 /// d = (b + t - 1) floordiv t
824 /// s = (b floordiv d) * sizeDivisor
825 /// v = b % d
826 /// u = d - v
827 ///
828 /// where the tile sizes are `s` and `s` + `sizeDivisor`, and the numbers of
829 /// the corresponding tiles are `u` and `v`, respectively. Alternatively,
830 ///
831 /// s * u + (s + sizeDivisor) * v == original size,
832 /// where s mod sizeDivisor = 0.
833 ///
834 /// Expects all values to be positive. In some cases with the target tile size
835 /// sufficiently close to the dimension shape and non-unit divisor, it is
836 /// impossible to compute such sizes. If `emitAssertion` is set, also emit the
837 /// assertion that size computation succeeded.
838 ///
839 /// Returns the specification consisting of both tile values and the number of
840 /// tiles of each size.
842 computeMultiTileSizes(OpBuilder &builder, LinalgOp op, unsigned dimension,
843  OpFoldResult targetSize, OpFoldResult divisor,
844  bool emitAssertions = true);
846 computeStaticMultiTileSizes(LinalgOp op, unsigned dimension, int64_t targetSize,
847  int64_t divisor);
848 
849 /// Rewrite a TilingInterface `op` to a tiled `scf.forall`, applying
850 /// tiling by `numThreads`.
851 /// If non-empty, the `mapping` is added as an attribute to the
852 /// resulting `scf.forall`.
853 /// Zero tile sizes indicate that the dimension is not tiled, and can be
854 /// thought of as tiling by the full size of data. It is the user's
855 /// responsibility to ensure that `numThreads` is a valid tiling specification
856 /// (i.e. that only tiles parallel dimensions, e.g. in the Linalg case).
860 };
862  TilingInterface op,
863  ArrayRef<OpFoldResult> numThreads,
864  std::optional<ArrayAttr> mapping);
865 
866 /// Same as `tileToForallOp`, but calculate the number of threads
867 /// required using the given tileSizes.
869 tileToForallOpUsingTileSizes(RewriterBase &builder, TilingInterface op,
870  ArrayRef<OpFoldResult> tileSizes,
871  std::optional<ArrayAttr> mapping);
872 
873 /// Transformation information returned after reduction tiling.
875  /// The partial reduction tiled op generated.
877  /// The final reduction operation merging all the partial reductions.
879  /// Initial values used for partial reductions.
881  /// The `scf.forall` operation that iterate over the tiles.
882  scf::ForallOp loops;
883 };
884 
885 /// Method to tile a reduction to parallel iterations computing partial
886 /// reductions. After the loop all the partial reduction are merged into a final
887 /// reduction. For example for the following sequence
888 ///
889 /// ```mlir
890 /// %0 = linalg.generic %in ["parallel", "reduction"]
891 /// : tensor<7x9xf32> -> tensor<7xf32>
892 /// ```
893 ///
894 /// into:
895 ///
896 /// ```mlir
897 /// %0 = linalg.fill ... : tensor<7x4xf32>
898 /// %1 = scf.forall (%iv) in (%c4) shared_outs(%arg0 = %0)
899 /// -> (tensor<7x4xf32>) {
900 /// %2 = tensor.extract_slice %arg3 : tensor<7x4xf32> to tensor<7xf32>
901 /// %3 = tensor.extract_slice %in : tensor<7x9xf32> -> tensor<7x?xf32>
902 /// %4 = linalg.generic %2, %3 ["parallel", "reduction"]
903 /// : tensor<7x?xf32> -> tensor<7xf32>
904 /// %5 = tensor.insert_slice %3, %arg0[0, %iv] : tensor<7x4xf32>
905 /// }
906 /// %6 = linalg.generic %1 ["parallel", "reduction"]
907 /// : tensor<7x4xf32> -> tensor<7xf32>
908 /// ```
910 tileReductionUsingForall(RewriterBase &b, PartialReductionOpInterface op,
911  ArrayRef<OpFoldResult> numThreads,
912  ArrayRef<OpFoldResult> tileSizes = {},
913  std::optional<ArrayAttr> mapping = std::nullopt);
914 
915 /// All indices returned by IndexOp should be invariant with respect to
916 /// tiling. Therefore, if an operation is tiled, we have to transform the
917 /// indices accordingly, i.e. offset them by the values of the corresponding
918 /// induction variables that are captured implicitly in the body of the op.
919 ///
920 /// Example. `linalg.generic` before tiling:
921 ///
922 /// #id_2d = (i, j) -> (i, j)
923 /// #pointwise_2d_trait = {
924 /// indexing_maps = [#id_2d, #id_2d],
925 /// iterator_types = ["parallel", "parallel"]
926 /// }
927 /// linalg.generic #pointwise_2d_trait %operand, %result {
928 /// ^bb0(%operand_in: f32, %result_in: f32):
929 /// %i = linalg.index 0 : index
930 /// %j = linalg.index 1 : index
931 /// <some operations that use %i, %j>
932 /// }: memref<50x100xf32>, memref<50x100xf32>
933 ///
934 /// After tiling pass with tiles sizes 10 and 25:
935 ///
936 /// #strided = (i, j)[s0, s1, s2] -> (i * s1 + s0 + j * s2)
937 ///
938 /// %c1 = arith.constant 1 : index
939 /// %c0 = arith.constant 0 : index
940 /// %c25 = arith.constant 25 : index
941 /// %c10 = arith.constant 10 : index
942 /// operand_dim_0 = dim %operand, 0 : memref<50x100xf32>
943 /// operand_dim_1 = dim %operand, 1 : memref<50x100xf32>
944 /// scf.for %k = %c0 to operand_dim_0 step %c10 {
945 /// scf.for %l = %c0 to operand_dim_1 step %c25 {
946 /// %4 = memref.subview %operand[%k, %l][%c10, %c25][%c1, %c1]
947 /// : memref<50x100xf32> to memref<?x?xf32, #strided>
948 /// %5 = memref.subview %result[%k, %l][%c10, %c25][%c1, %c1]
949 /// : memref<50x100xf32> to memref<?x?xf32, #strided>
950 /// linalg.generic pointwise_2d_trait %4, %5 {
951 /// ^bb0(%operand_in: f32, %result_in: f32):
952 /// %i = linalg.index 0 : index
953 /// %j = linalg.index 1 : index
954 /// // Indices `k` and `l` are implicitly captured in the body.
955 /// %transformed_i = arith.addi %i, %k : index // index `i` is offset by
956 /// %k %transformed_j = arith.addi %j, %l : index // index `j` is offset
957 /// by %l
958 /// // Every use of %i, %j is replaced with %transformed_i,
959 /// %transformed_j <some operations that use %transformed_i,
960 /// %transformed_j>
961 /// }: memref<?x?xf32, #strided>, memref<?x?xf32, #strided>
962 /// }
963 /// }
964 ///
965 /// TODO: Investigate whether mixing implicit and explicit indices
966 /// does not lead to losing information.
967 void transformIndexOps(RewriterBase &b, LinalgOp op,
969  const LoopIndexToRangeIndexMap &loopIndexToRangeIndex);
970 
971 /// Apply transformation to split the single linalg op reduction into a
972 /// parallel and reduction dimension. Then create a new linalg.generic op
973 /// doing the rest of the reduction. Return the new linalg op with an extra
974 /// parallel dimension or failure if the transformation didn't happen.
975 ///
976 /// Example:
977 /// ```
978 /// %r = linalg.generic {indexing_maps = [affine_map<(d0) -> (d0)>,
979 /// affine_map<(d0) -> ()>],
980 /// iterator_types = ["reduction"]}
981 /// ins(%in : tensor<32xf32>)
982 /// outs(%out : tensor<f32>) {
983 /// ^bb0(%arg1: f32, %arg2: f32):
984 /// %y = arith.addf %arg1, %arg2 : f32
985 /// linalg.yield %y : f32
986 /// } -> tensor<f32>
987 /// ```
988 /// To:
989 /// ```
990 /// %cst = arith.constant 0.000000e+00 : f32
991 /// %0 = tensor.expand_shape %in [[0, 1]] : tensor<32xf32> into
992 /// tensor<4x8xf32> %1 = tensor.empty [4] : tensor<4xf32> %2 = linalg.fill
993 /// ins(%cst : f32) outs(%1 : tensor<4xf32>) -> tensor<4xf32> %3 =
994 /// linalg.generic {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>,
995 /// affine_map<(d0, d1) -> (d0)>],
996 /// iterator_types = ["parallel", "reduction"]}
997 /// ins(%0 : tensor<4x8xf32>) outs(%2 : tensor<4xf32>) {
998 /// ^bb0(%arg3: f32, %arg5: f32):
999 /// %5 = arith.addf %arg3, %arg4 : f32
1000 /// linalg.yield %5 : f32
1001 /// } -> tensor<4xf32>
1002 /// %r = linalg.generic {indexing_maps = [affine_map<(d0) -> (d0)>,
1003 /// affine_map<(d0) -> ()>],
1004 /// iterator_types = ["reduction"]}
1005 /// ins(%3 : tensor<4xf32>) outs(%out : tensor<f32>) {
1006 /// ^bb0(%arg3: f32, %arg4: f32):
1007 /// %5 = arith.addf %arg3, %arg4 : f32
1008 /// linalg.yield %5 : f32
1009 /// } -> tensor<f32>
1010 /// ```
1013  FillOp fillOp;
1014  LinalgOp splitLinalgOp;
1016 };
1018 splitReduction(RewriterBase &b, LinalgOp op,
1019  const ControlSplitReductionFn &controlSplitReductionFn,
1020  bool useAlloc = false);
1021 
1022 /// Scaling-based implementation of the split reduction transformation.
1023 /// Instead of introducing an ExpandShapeOp, this rewrites a reduction
1024 /// dimension `k` into `k * scale + kk`.
1025 ///
1026 /// Example:
1027 /// ```
1028 /// %0 = linalg.matmul ins(%A, %B: tensor<16x256xf32>, tensor<256x32xf32>)
1029 /// outs(%C: tensor<16x32xf32>) -> tensor<16x32xf32>
1030 /// ```
1031 ///
1032 /// Is transformed to:
1033 ///
1034 /// ```
1035 /// #map0 = affine_map<(d0, d1, d2, d3) -> (d0, d2 * 4 + d3)>
1036 /// #map1 = affine_map<(d0, d1, d2, d3) -> (d2 * 4 + d3, d1)>
1037 /// #map2 = affine_map<(d0, d1, d2, d3) -> (d2, d3)>
1038 /// #map3 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>
1039 /// #map4 = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
1040 /// #map5 = affine_map<(d0, d1, d2) -> (d0, d1)>
1041 /// %0 = tensor.empty [16, 32, 64] : tensor<16x32x64xf32>
1042 /// %cst = arith.constant 0.000000e+00 : f32
1043 /// %1 = linalg.fill ins(%cst : f32) outs(%0 : tensor<16x32x64xf32>) ->
1044 /// tensor<16x32x64xf32>
1045 /// %2 = tensor.empty [64, 4] : tensor<64x4xi1>
1046 ///
1047 /// %3 = linalg.generic {indexing_maps = [#map0, #map1, #map2, #map3],
1048 /// iterator_types = ["parallel", "parallel", "parallel", "reduction"]}
1049 /// ins(%A, %B, %2 : tensor<16x256xf32>, tensor<256x32xf32>,
1050 /// tensor<64x4xi1>)
1051 /// outs(%1 : tensor<16x32x64xf32>) {
1052 /// ^bb0(%arg3: f32, %arg4: f32, %arg5: i1, %arg6: f32):
1053 /// %5 = arith.mulf %arg3, %arg4 : f32
1054 /// %6 = arith.addf %arg6, %5 : f32
1055 /// linalg.yield %6 : f32
1056 /// } -> tensor<16x32x64xf32>
1057 ///
1058 /// %4 = linalg.generic {indexing_maps = [#map4, #map5],
1059 /// iterator_types = ["parallel", "parallel", "reduction"]}
1060 // ins(%3 : tensor<16x32x64xf32>)
1061 /// outs(%C : tensor<16x32xf32>) {
1062 /// ^bb0(%arg3: f32, %arg4: f32):
1063 /// %5 = arith.addf %arg3, %arg4 : f32
1064 /// linalg.yield %5 : f32
1065 /// } -> tensor<16x32xf32>
1066 ///
1067 /// return %4 : tensor<16x32xf32>
1068 /// ```
1070 splitReductionByScaling(RewriterBase &b, LinalgOp op,
1071  const ControlSplitReductionFn &controlSplitReductionFn,
1072  bool useAlloc = false);
1073 
1074 /// Return `true` if a given sequence of dimensions are contiguous in the
1075 /// range of the specified indexing map.
1077 /// Return `true` if all sequences of dimensions specified in `dimSequences` are
1078 /// contiguous in all the ranges of the `maps`.
1080  ArrayRef<ReassociationIndices> dimSequences);
1081 
1084  LinalgOp collapsedOp;
1085 };
1086 
1087 /// Collapses dimensions of linalg.generic/linalg.copy operation. A precondition
1088 /// to calling this method is that for each list in `foldedIterationDim`, the
1089 /// sequence of dimensions is contiguous in domains of all `indexing_maps` of
1090 /// the `linalgOp`. This can be checked using `areDimSequencePreserved` method.
1091 /// When valid, the method also collapses the operands of the op. Returns
1092 /// replacement values of the results of the original `linalgOp` by inserting
1093 /// reshapes to get back values of compatible types.
1095 collapseOpIterationDims(LinalgOp op,
1096  ArrayRef<ReassociationIndices> foldedIterationDims,
1097  RewriterBase &rewriter);
1098 
1100  tensor::PadOp padOp;
1101  tensor::ExpandShapeOp expandShapeOp;
1102  linalg::TransposeOp transposeOp;
1103 };
1104 
1105 /// Rewrite pack as pad + reshape + transpose.
1107  tensor::PackOp packOp);
1108 
1110  tensor::EmptyOp emptyOp;
1111  linalg::TransposeOp transposeOp;
1112  tensor::CollapseShapeOp collapseShapeOp;
1113  tensor::ExtractSliceOp extractSliceOp;
1114 };
1115 
1116 /// Rewrite pack as empty + transpose + reshape + extract_slice.
1118  tensor::UnPackOp unPackOp);
1119 
1120 /// Struct to hold the result of a `pack` call.
1121 struct PackResult {
1123  linalg::LinalgOp packedLinalgOp;
1125 };
1126 /// Implement packing of a single LinalgOp by `packedSizes`.
1127 /// There must be one packedSizes entry per `linalgOp` iterator.
1128 /// Return the packed Linalg op on success, failure otherwise.
1129 FailureOr<PackResult> pack(RewriterBase &rewriter, linalg::LinalgOp linalgOp,
1130  ArrayRef<OpFoldResult> packedSizes);
1131 
1132 /// Struct to hold the result of a `packTranspose` call.
1134  tensor::PackOp transposedPackOp;
1135  linalg::LinalgOp transposedLinalgOp;
1136  tensor::UnPackOp transposedUnPackOp;
1137 };
1138 /// Transpose a single PackOp -> LinalgOp -> UnPackOp chain and return the
1139 /// transposed PackOp -> LinalgOp -> UnPackOp chain after replacements.
1140 /// Return failure if either:
1141 /// 1. the `packOp` does not have the `linalgOp` as its unique use.
1142 /// 2. the `maybeUnPackOp`, if specified must be a consumer of the result tied
1143 /// to the unique `packOp` use.
1144 /// 3. `outerPerm` (resp. `innerPerm`) must be valid permutations of
1145 /// `packOp.getOuterDimsPerm` (resp. `packOp.getInnerDimsPerm`) or empty.
1147 packTranspose(RewriterBase &rewriter, tensor::PackOp packOp,
1148  linalg::LinalgOp linalgOp, tensor::UnPackOp maybeUnPackOp,
1149  ArrayRef<int64_t> outerPerm, ArrayRef<int64_t> innerPerm);
1150 
1151 /// Pack a LinalgOp by greedily inferring matmul dimensions (m, n, k) where m
1152 /// and n are proper parallel dimensions and k is a proper reduction
1153 /// dimension. Packing occurs by rewriting the op as a linalg.generic and
1154 /// calling linalg::pack by `mnkPackedSizes`. The order of the packed
1155 /// dimensions is customizable: the `mnkOrder` is a permutation of {0, 1, 2}
1156 /// to reorder {m, n, k} into one of the 8 possible forms. The outer
1157 /// dimensions of the operands are not permuted at this time, this is left for
1158 /// future work.
1160 packMatmulGreedily(RewriterBase &rewriter, LinalgOp linalgOp,
1161  ArrayRef<OpFoldResult> mnkPackedSizes,
1162  ArrayRef<int64_t> mnkPaddedSizesNextMultipleOf,
1163  ArrayRef<int64_t> mnkOrder);
1164 
1166  /// Minor block factors (mb, nb, kb) for packing relayout where mb, mn are
1167  /// the parallel dimensions and kb is the reduction dimension.
1169 
1170  /// If true, allows packing of dimensions that only partially fit into the
1171  /// block factors.
1172  bool allowPadding = true;
1173 
1174  /// Next multiples of the packing sizes.
1176 
1177  /// Permutation of matmul (M, N, K) dimensions order.
1179 
1180  /// Transpose LHS outer block layout [MB][KB] -> [KB][MB].
1182 
1183  /// Transpose LHS inner block layout [mb][kb] -> [kb][mb].
1185 
1186  /// Transpose RHS outer block layout [KB][NB] -> [NB][KB].
1188 
1189  /// Transpose RHS inner block layout [kb][nb] -> [nb][kb].
1191 };
1192 
1193 /// Function type which is used to control matmul packing.
1194 /// It is expected to return valid packing configuration for each operation.
1195 /// Lack of packing options indicates that no valid configuration could be
1196 /// assigned and the operation will not be packed.
1198  std::function<std::optional<BlockPackMatmulOptions>(linalg::LinalgOp)>;
1199 
1200 /// Pack a matmul operation into blocked 4D layout.
1201 ///
1202 /// Relayout a matmul operation into blocked layout with two levels of
1203 /// subdivision:
1204 /// - major 2D blocks - outer dimensions, consist of minor blocks
1205 /// - minor 2D blocks - inner dimensions, consist of scalar elements
1206 ///
1207 /// A 2D matmul MxNxK gets reshaped into blocked 4D representation
1208 /// as: [MB][NB][mb][nb] += [MB][KB][mb][kb] * [NB][KB][nb][kb]
1209 /// where the (MB, NB, KB) dimensions represent the major blocks,
1210 /// and the (mb, nb, kb) are the minor blocks of their respective
1211 /// original 2D dimensions (M, N, K).
1212 ///
1213 /// Depending on the initial operands' data layout and the specified
1214 /// packing options, the major blocks dimensions might get transposed
1215 /// e.g., [MB][KB] -> [KB][MB]. The minor blocks can also be transposed
1216 /// e.g., [mb][kb] -> [kb][mb].
1217 /// Any present batch dimensions remain unchanged.
1218 /// The final result is unpacked back to the original shape.
1219 ///
1220 /// Return failure if no valid packing options are provided.
1222 blockPackMatmul(RewriterBase &rewriter, linalg::LinalgOp linalgOp,
1223  const ControlBlockPackMatmulFn &controlPackMatmul);
1224 
1225 /// Rewrite tensor.from_elements to linalg.generic.
1228  tensor::FromElementsOp fromElementsOp);
1229 
1230 /// Rewrite tensor.generate to linalg.generic.
1233  tensor::GenerateOp generateOp);
1234 
1235 /// Rewrite tensor.pad to linalg.generic + tensor.insert_slice.
1237  tensor::PadOp padOp);
1238 
1239 /// Convert linalg.conv_2d_nhwc_hwcf into linalg.generic (for img2col packing)
1240 /// and linalg.matmul.
1241 ///
1242 /// A convolution operation can be written as a matrix-matrix multiplication by
1243 /// unfolding the cross-correlation between input and filter and explicitly copy
1244 /// overlapped sliding window inputs.
1245 ///
1246 /// Consider 2D input X with single channel input and output and 2x2 filter W:
1247 /// [x(0, 0) , x(0, 1) , ..., x(0, n) ]
1248 /// [x(1, 0) , x(1, 1) , ..., x(1, n) ]
1249 /// [. , . ,. , . ] [w(0, 0), w(0, 1)]
1250 /// [. , . , . , . ] (conv) [w(1, 0), w(1, 1)]
1251 /// [. , . , ., . ]
1252 /// [x(n-1, 0), x(n-1, 1), ..., x(n-1, n-1)]
1253 ///
1254 /// The packed input data (img2col) is a matrix with |rows| = output spatial
1255 /// size, |columns| = filter spatial size. To compute the output Y(i, j) we need
1256 /// to calculate the dot product between filter window at input X(x, y)) and the
1257 /// filter which will look like the following where r.h.s is the img2col matrix
1258 /// and l.h.s is the flattened filter:
1259 ///
1260 /// [x(0,0), x(0,1), x(1,0), x(1,1)]
1261 /// [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)]
1262 /// [x(0,1), x(1,1), x(0,2), x(1,2)]
1263 /// [ . , . , . , . ]
1264 ///
1265 /// In general for 2D case with (N, H, W, C) input and (Kh, Kw, C, D) filter
1266 /// and output (N, Ho, Wo, D) the convolution is the following matrix-matrix
1267 /// multiplication (Ho x Wo, Kh x Kw x C) * (Kh x Kw x C, D) for each input in
1268 /// the N input. For the case where N > 1 its a batched matrix-matrix
1269 /// multiplication.
1270 ///
1271 /// On success, return both the operation that produces the img2col tensor and
1272 /// the final operation of the sequence that replaces the original convolution.
1274 rewriteInIm2Col(RewriterBase &rewriter, linalg::Conv2DNhwcHwcfOp convOp);
1275 
1276 /// Same as the above but for Fhwc channel orderings in the filter. In this case
1277 /// the matrix multiplication is actually a row-wise dot-product rather than a
1278 /// row-column dot-product. This is to avoid transposing the filter matrix which
1279 /// would be required for a regular matrix multiplication to produce the correct
1280 /// output dimensions.
1282 rewriteInIm2Col(RewriterBase &rewriter, linalg::Conv2DNhwcFhwcOp convOp);
1283 
1284 /// Similar to rewriteInIm2Col with linalg::Conv2DNhwcHwcfOp except there is no
1285 /// reduction among the input channels so each convolution can be a
1286 /// matrix-vector product and by transposing both input filter so channels are
1287 /// outer most the computation is a batched matrix-vector product.
1289 rewriteInIm2Col(RewriterBase &rewriter,
1290  linalg::DepthwiseConv2DNhwcHwcOp convOp);
1291 
1292 /// Similar to rewriteInIm2Col with linalg::Conv2DNhwcHwcfOp except because the
1293 /// channels are to the left of the image shape dimensions, the position of the
1294 /// contraction dimension in the resulting matmul is reversed. This swaps the
1295 /// LHS and RHS of the matmul when compared with nhwc (i.e. (D, C x Kh x Kw) *
1296 /// (C x Kh x Kw, Ho x Wo))
1298 rewriteInIm2Col(RewriterBase &rewriter, linalg::Conv2DNchwFchwOp convOp);
1299 
1300 /// Convert linalg.conv_2d_nhwc_fhwc(_q) to linalg.conv_2d_nhwc_hwcf(_q) by
1301 /// materializing transpose.
1303  linalg::Conv2DNhwcFhwcOp op);
1305  linalg::Conv2DNhwcFhwcQOp op);
1306 
1307 /// Convert Linalg matmul ops to transposed variants.
1309  linalg::MatmulOp op,
1310  bool transposeLHS = true);
1312  linalg::BatchMatmulOp op,
1313  bool transposeLHS = true);
1314 
1315 //===----------------------------------------------------------------------===//
1316 // Rewrite patterns wrapping transformations.
1317 // TODO: every single such pattern should be a close to noop wrapper around a
1318 // functional-stye API call.
1319 //===----------------------------------------------------------------------===//
1320 
1321 /// Rewrites 2-D convolution ops with size-1 window dimensions into 1-D
1322 /// convolution ops.
1323 template <typename Conv2DOp, typename Conv1DOp>
1325  : public OpRewritePattern<Conv2DOp> {
1327 
1329  PatternRewriter &rewriter) const;
1330 
1332  PatternRewriter &rewriter) const override {
1333  return returningMatchAndRewrite(convOp, rewriter);
1334  }
1335 };
1336 
1337 extern template struct DownscaleSizeOneWindowed2DConvolution<Conv2DNhwcHwcfOp,
1338  Conv1DNwcWcfOp>;
1339 extern template struct DownscaleSizeOneWindowed2DConvolution<Conv2DNchwFchwOp,
1340  Conv1DNcwFcwOp>;
1341 
1342 /// Rewrites 2-D depthwise convolution ops with size-1 (w, kw) or (h, kh)
1343 /// dimensions into 1-D depthwise convolution ops.
1345  : public OpRewritePattern<DepthwiseConv2DNhwcHwcOp> {
1347  PatternBenefit benefit = 1)
1348  : OpRewritePattern<DepthwiseConv2DNhwcHwcOp>(context, benefit) {}
1349 
1351  returningMatchAndRewrite(DepthwiseConv2DNhwcHwcOp convOp,
1352  PatternRewriter &rewriter) const;
1353 
1354  LogicalResult matchAndRewrite(DepthwiseConv2DNhwcHwcOp convOp,
1355  PatternRewriter &rewriter) const override {
1356  return returningMatchAndRewrite(convOp, rewriter);
1357  }
1358 };
1359 
1360 struct DownscaleConv2DOp final : public OpRewritePattern<Conv2DOp> {
1362  : OpRewritePattern<Conv2DOp>(context, benefit) {}
1363 
1365  PatternRewriter &rewriter) const;
1366 
1368  PatternRewriter &rewriter) const override {
1369  return returningMatchAndRewrite(convOp, rewriter);
1370  }
1371 };
1372 
1373 ///
1374 /// Linalg generalization pattern.
1375 ///
1376 /// Apply the `generalization` transformation as a pattern.
1377 /// See `generalization` for more details.
1378 //
1379 // TODO: Automatic default pattern class that just unwraps a function
1380 // returning FailureOr<GenericOp>.
1382  : public OpInterfaceRewritePattern<LinalgOp> {
1384 
1385  /// `matchAndRewrite` implementation that returns the significant
1386  /// transformed pieces of IR.
1388  returningMatchAndRewrite(LinalgOp op, PatternRewriter &rewriter) const {
1389  return generalizeNamedOp(rewriter, op);
1390  }
1391 
1393  PatternRewriter &rewriter) const override {
1394  return returningMatchAndRewrite(op, rewriter);
1395  }
1396 };
1397 
1398 /// Vectorization pattern for memref::CopyOp.
1399 struct CopyVectorizationPattern : public OpRewritePattern<memref::CopyOp> {
1401 
1402  LogicalResult matchAndRewrite(memref::CopyOp copyOp,
1403  PatternRewriter &rewriter) const override;
1404 };
1405 
1407  std::function<LogicalResult(RewriterBase &, tensor::PadOp, Value)>;
1408 
1409 /// Rewrite a tensor::PadOp into a sequence of EmptyOp, FillOp and
1410 /// InsertSliceOp. For now, only constant padding values are supported.
1411 /// `OptimizeCopyFn` can be used to customize copying step optimization.
1412 struct GeneralizePadOpPattern : public OpRewritePattern<tensor::PadOp> {
1414  OptimizeCopyFn optimizeCopyFn = nullptr,
1415  PatternBenefit benefit = 1)
1416  : OpRewritePattern<tensor::PadOp>(context, benefit),
1417  optimizeCopyFn(std::move(optimizeCopyFn)) {}
1418  LogicalResult matchAndRewrite(tensor::PadOp padOp,
1419  PatternRewriter &rewriter) const override;
1420 
1421 protected:
1423  Value createFillOrGenerateOp(RewriterBase &rewriter, tensor::PadOp padOp,
1424  Value dest,
1425  const SmallVector<Value> &dynSizes) const;
1426 };
1427 
1428 /// Rewrites a tensor::PackOp into a sequence of tensor.pad + linalg.transpose +
1429 /// tensor.insert_slice ops, where the tensor::PackOp has outer dims being all
1430 /// 1s.
1432  : public OpRewritePattern<tensor::PackOp> {
1434  LogicalResult matchAndRewrite(tensor::PackOp packOp,
1435  PatternRewriter &rewriter) const override;
1436 };
1437 
1438 /// Rewrites a tensor::UnPackOp into a sequence of rank-reduced extract_slice op
1439 /// + transpose op + insert_slice op, where the tensor::UnPackOp has outer dims
1440 /// being all 1s.
1442  : public OpRewritePattern<tensor::UnPackOp> {
1444  LogicalResult matchAndRewrite(tensor::UnPackOp unpackOp,
1445  PatternRewriter &rewriter) const override;
1446 };
1447 
1448 /// Match and rewrite for the pattern:
1449 /// ```
1450 /// %alloc = ...
1451 /// [optional] %view = memref.view %alloc ...
1452 /// %subView = subview %allocOrView ...
1453 /// [optional] linalg.fill(%allocOrView, %cst) ...
1454 /// ...
1455 /// memref.copy(%in, %subView) ...
1456 /// vector.transfer_read %allocOrView[...], %cst ...
1457 /// ```
1458 /// into
1459 /// ```
1460 /// [unchanged] %alloc = ...
1461 /// [unchanged] [optional] %view = memref.view %alloc ...
1462 /// [unchanged] [unchanged] %subView = subview %allocOrView ...
1463 /// ...
1464 /// vector.transfer_read %in[...], %cst ...
1465 /// ```
1466 /// Where there is no interleaved use between memref.copy and transfer_read as
1467 /// well as no interleaved use between linalg.fill and memref.copy (if
1468 /// linalg.fill is specified).
1469 /// This is a custom rewrite to forward partial reads (with optional fills) to
1470 /// vector.transfer_read.
1472  : public OpRewritePattern<vector::TransferReadOp> {
1474 
1475  LogicalResult matchAndRewrite(vector::TransferReadOp xferOp,
1476  PatternRewriter &rewriter) const override;
1477 };
1478 
1479 /// Match and rewrite for the pattern:
1480 /// ```
1481 /// %alloc = ...
1482 /// [optional] %view = memref.view %alloc ...
1483 /// %subView = subview %allocOrView...
1484 /// ...
1485 /// vector.transfer_write %..., %allocOrView[...]
1486 /// memref.copy(%subView, %out)
1487 /// ```
1488 /// into
1489 /// ```
1490 /// [unchanged] %alloc = ...
1491 /// [unchanged] [optional] %view = memref.view %alloc ...
1492 /// [unchanged] %subView = subview %allocOrView...
1493 /// ...
1494 /// vector.transfer_write %..., %out[...]
1495 /// ```
1496 /// Where there is no interleaved use between transfer_write and memref.copy.
1497 /// This is a custom rewrite to forward partial writes to
1498 /// vector.transfer_write.
1500  : public OpRewritePattern<vector::TransferWriteOp> {
1502 
1503  LogicalResult matchAndRewrite(vector::TransferWriteOp xferOp,
1504  PatternRewriter &rewriter) const override;
1505 };
1506 
1507 /// Rewrite extract_slice(tensor.pad(x)) into tensor.pad(extract_slice(x)).
1509  : public OpRewritePattern<tensor::ExtractSliceOp> {
1510  /// A function to control pattern application and rewrite logic.
1511  ///
1512  /// The function will be given the slice op and should return:
1513  /// - std::nullopt: to fail the match and not apply the pattern;
1514  /// - true: to apply the pattern with zero slice guard;
1515  /// - false: to apply the pattern without zero slice guard.
1516  ///
1517  /// See the documentation for tensor::bubbleUpPadSlice regarding zero slice
1518  /// guard.
1519  using ControlFn = std::function<std::optional<bool>(tensor::ExtractSliceOp)>;
1520 
1522  ControlFn controlFn = nullptr,
1523  PatternBenefit benefit = 1)
1524  : OpRewritePattern(context, benefit), controlFn(std::move(controlFn)) {}
1525 
1526  LogicalResult matchAndRewrite(tensor::ExtractSliceOp sliceOp,
1527  PatternRewriter &rewriter) const override;
1528 
1529 private:
1530  ControlFn controlFn;
1531 };
1532 
1533 //===----------------------------------------------------------------------===//
1534 // Populate functions.
1535 //===----------------------------------------------------------------------===//
1536 
1537 /// Canonicalization patterns relevant to apply after tiling patterns. These
1538 /// are applied automatically by the tiling pass but need to be applied
1539 /// manually when tiling is called programmatically.
1542 
1543 /// Linalg generalization patterns
1544 
1545 /// Populates `patterns` with patterns to convert spec-generated named ops to
1546 /// linalg.generic ops.
1548 
1549 /// Linalg decompose convolutions patterns
1550 
1551 /// Populates patterns to decompose high-D convolution ops into low-D ones.
1552 /// This is a step in progressive lowering for convolution ops, afterwards we
1553 /// can vectorize the low-D convolution ops.
1555  PatternBenefit benefit = 1);
1556 
1557 /// Populates patterns to transform linalg.conv_2d_xxx operations into
1558 /// linalg.generic (for img2col packing) and linalg.matmul.
1559 /// \see rewriteInIm2Col for more details.
1561 
1562 /// Populates `patterns` with patterns that vectorize tensor.pad.
1563 /// These patterns are meant to apply in a complementary fashion. Benefits
1564 /// are used to encode a certain ordering of pattern application. To avoid
1565 /// scattering magic constants throughout the code base, the patterns must be
1566 /// added with this function. `baseBenefit` can be used to offset the benefit
1567 /// of all tensor::PadOp vectorization patterns by a certain value.
1569  PatternBenefit baseBenefit = 1);
1570 
1571 /// Populate patterns for splitting a `LinalgOp` with multiple statements within
1572 /// its payload into multiple `GenericOp` that have a single statement.
1573 /// The option `removeDeadArgsAndResults` adds patterns to remove dead arguments
1574 /// and results from the generated decomposed ops. This is default `true` since
1575 /// the core decomposition patterns relies on these clean up patterns. It is set
1576 /// to false only for testing purposes.
1578  bool removeDeadArgsAndResults = true);
1579 
1580 /// Populate patterns that convert non-destination-style ops to destination
1581 /// style ops.
1583 
1584 /// Populate patterns for vectorizing low-D convolution ops. This is a step in
1585 /// progressive lowering for convolution ops, it assume high-D convolution ops
1586 /// were decomposed previously.
1588  PatternBenefit benefit = 1);
1589 
1590 /// Populate patterns that convert `ElementwiseMappable` ops to linalg
1591 /// parallel loops.
1593 
1594 /// Populate patterns that are only useful in the context of sparse tensors.
1596 
1597 /// Function type which is used to control when to stop fusion. It is expected
1598 /// that OpOperand is not modified in the callback. The OpOperand is not marked
1599 /// as const to allow callers to use non-const methods.
1600 using ControlFusionFn = std::function<bool(OpOperand *fusedOperand)>;
1601 
1602 /// Patterns for fusing linalg operation on tensors.
1603 
1604 /// Pattern to fuse `linalg.generic` -> `linalg.generic` operations
1605 /// when both operations are fusable elementwise operations.
1607  RewritePatternSet &patterns,
1608  const ControlFusionFn &controlElementwiseOpFusion);
1609 
1610 /// Function type which is used to control propagation of tensor.pack/unpack
1611 /// ops.
1612 using ControlPropagationFn = std::function<bool(Operation *op)>;
1613 
1614 /// Patterns to bubble up or down data layout ops across other operations.
1616  RewritePatternSet &patterns,
1617  const ControlPropagationFn &controlPackUnPackPropagation);
1618 
1619 /// Pattern to remove dead operands and results of `linalg.generic` operations.
1620 /// This is effectively DCE for a linalg op.
1622 
1623 /// Patterns to promote inputs to outputs and remove unused inputs of
1624 /// `linalg.generic` ops.
1626 
1627 /// Function type to control generic op dimension collapsing. It is expected
1628 /// to return an array of `ReassociationIndices` representing dimensions that
1629 /// should be merged.
1631  std::function<SmallVector<ReassociationIndices>(linalg::LinalgOp)>;
1632 
1633 /// Pattern to collapse dimensions in a linalg.generic op. This will collapse
1634 /// tensor operands when needed and expand back the result tensors.
1636  RewritePatternSet &patterns,
1637  const GetCollapsableDimensionsFn &controlCollapseDimensions);
1638 
1639 /// Patterns to fold an expanding (collapsing) tensor_reshape operation with its
1640 /// producer (consumer) generic operation by expanding the dimensionality of the
1641 /// loop in the generic op.
1643  RewritePatternSet &patterns, const ControlFusionFn &controlFoldingReshapes);
1644 
1645 /// Patterns to fold an expanding tensor.expand_shape operation with its
1646 /// producer generic operation by collapsing the dimensions of the generic op.
1648  RewritePatternSet &patterns, const ControlFusionFn &controlFoldingReshapes);
1649 
1650 /// Patterns to constant fold Linalg operations.
1652  const ControlFusionFn &controlFn);
1653 
1654 /// Pattern to fuse a `tensor.pad` operation with the producer of its source,
1655 /// if the producer is a `linalg` operation with all parallel iterator types.
1657  RewritePatternSet &patterns);
1658 
1659 /// Patterns to convert from one named op to another. These can be seen as
1660 /// canonicalizations of named ops into another named op.
1662 
1663 /// Patterns to fold unit-extent dimensions in operands/results of linalg ops on
1664 /// tensors via reassociative reshape ops.
1667 
1668 /// A pattern that converts init operands to input operands.
1670 
1671 /// Patterns that are used to inline constant operands into linalg generic ops.
1673 
1674 /// Patterns that are used to bubble up extract slice op above linalg op.
1676 
1677 /// Adds patterns that waps tensor.extract_slice(linalg.fill(%cst, %init)) into
1678 /// linalg.fill(%cst, tensor.extract_slice(%init)).
1680 
1681 /// Patterns to apply `splitReduction` below.
1683  RewritePatternSet &patterns,
1684  const ControlSplitReductionFn &controlSplitReductionFn,
1685  bool useAlloc = false);
1686 
1687 /// Patterns to convert Linalg matmul ops to transposed variants.
1689  bool transposeLHS = true);
1690 
1691 /// Patterns to block pack Linalg matmul ops.
1693  const ControlBlockPackMatmulFn &controlFn);
1694 
1695 } // namespace linalg
1696 } // namespace mlir
1697 
1698 #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:47
Attributes are known-constant values of operations.
Definition: Attributes.h:25
The main mechanism for performing data layout queries.
This class provides support for representing a failure result, or a valid value of type T.
Definition: LogicalResult.h:78
This class defines the main interface for locations in MLIR and acts as a non-nullable wrapper around...
Definition: Location.h:63
MLIRContext is the top-level object for a collection of MLIR operations.
Definition: MLIRContext.h:60
This class helps build Operations.
Definition: Builders.h:209
This class represents a single result from folding an operation.
Definition: OpDefinition.h:268
This class represents an operand of an operation.
Definition: Value.h:267
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:785
This class coordinates the application of a rewrite on a set of IR, providing a way for clients to tr...
Definition: PatternMatch.h:400
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.
FailureOr< GenericOp > generalizeNamedOp(RewriterBase &rewriter, LinalgOp namedOp)
Create a GenericOp from the given named operation namedOp and replace namedOp.
void populateTransposeMatmulPatterns(RewritePatternSet &patterns, bool transposeLHS=true)
Patterns to convert Linalg matmul ops to transposed variants.
LogicalResult rewriteAsPaddedOp(RewriterBase &rewriter, LinalgOp opToPad, const LinalgPaddingOptions &options, LinalgOp &paddedOp, SmallVector< Value > &replacements, SmallVector< tensor::PadOp > &padOps)
Pad the iterator dimensions paddingDimensions of all opToPad operands to a static bounding box.
Definition: Padding.cpp:153
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...
FailureOr< ForallTilingResult > tileToForallOpUsingTileSizes(RewriterBase &builder, TilingInterface op, ArrayRef< OpFoldResult > tileSizes, std::optional< ArrayAttr > mapping)
Same as tileToForallOp, but calculate the number of threads required using the given tileSizes.
Definition: Tiling.cpp:467
FailureOr< LowerUnPackOpResult > lowerUnPack(RewriterBase &rewriter, tensor::UnPackOp unPackOp)
Rewrite pack as empty + transpose + reshape + extract_slice.
Definition: Transforms.cpp:357
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:78
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:340
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...
DenseMap< int, int > LoopIndexToRangeIndexMap
Creates a number of ranges equal to the number of non-zero in tileSizes.
Definition: Transforms.h:790
std::optional< Value > allocateWorkgroupMemory(OpBuilder &builder, memref::SubViewOp subview, ArrayRef< Value > sizeBounds, DataLayout &)
Allocate the subview in the GPU workgroup memory.
Definition: Promotion.cpp:470
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:1600
bool isDimSequencePreserved(AffineMap map, ReassociationIndicesRef dimSequence)
Return true if a given sequence of dimensions are contiguous in the range of the specified indexing m...
std::function< bool(Operation *op)> ControlPropagationFn
Function type which is used to control propagation of tensor.pack/unpack ops.
Definition: Transforms.h:1612
FailureOr< LinalgOp > specializeGenericOp(RewriterBase &rewriter, GenericOp genericOp)
Create a namedOp from the given GenericOp and replace the GenericOp.
Definition: Specialize.cpp:61
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:103
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:345
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:912
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:511
void populateSparseTensorRewriting(RewritePatternSet &patterns)
Populate patterns that are only useful in the context of sparse tensors.
FailureOr< ElementwiseOpFusionResult > fuseElementwiseOps(RewriterBase &rewriter, OpOperand *fusedOperand)
FailureOr< PromotionInfo > promoteSubviewAsNewBuffer(OpBuilder &b, Location loc, memref::SubViewOp subView, const AllocBufferCallbackFn &allocationFn, DataLayout &layout)
Definition: Promotion.cpp:238
std::optional< Value > allocateGPUPrivateMemory(OpBuilder &builder, memref::SubViewOp subview, ArrayRef< Value > sizeBounds, DataLayout &)
Allocate the subview in the GPU private memory.
Definition: Promotion.cpp:495
FailureOr< Operation * > rewriteInDestinationPassingStyle(RewriterBase &rewriter, tensor::FromElementsOp fromElementsOp)
Rewrite tensor.from_elements to linalg.generic.
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:75
FailureOr< ForallTilingResult > tileToForallOp(RewriterBase &builder, TilingInterface op, ArrayRef< OpFoldResult > numThreads, std::optional< ArrayAttr > mapping)
Definition: Tiling.cpp:458
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:486
std::function< SmallVector< Value, 4 >(OpBuilder &, Operation *)> TileSizeComputationFunction
Definition: Transforms.h:187
std::function< LogicalResult(RewriterBase &, tensor::PadOp, Value)> OptimizeCopyFn
Definition: Transforms.h:1407
FailureOr< Value > hoistPaddingOnTensors(RewriterBase &rewriter, tensor::PadOp opToHoist, int64_t numLoops, ArrayRef< int64_t > transposeVector, tensor::PadOp &hoistedOp, SmallVectorImpl< GenericOp > &transposeOps)
Mechanically hoist padding operations on tensors by numLoops into a new, generally larger tensor.
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:368
LogicalResult vectorize(RewriterBase &rewriter, Operation *op, ArrayRef< int64_t > inputVectorSizes={}, ArrayRef< bool > inputScalableVecDims={}, bool vectorizeNDExtract=false, bool flatten1DDepthwiseConv=false)
Emit a suitable vector form for an operation.
std::tuple< SmallVector< Range, 4 >, LoopIndexToRangeIndexMap > makeTiledLoopRanges(RewriterBase &b, Location loc, AffineMap map, ArrayRef< OpFoldResult > allShapeSizes, ArrayRef< OpFoldResult > allTileSizes)
Definition: Tiling.cpp:49
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:399
LogicalResult copyToGPUPrivateMemory(OpBuilder &b, Value src, Value dst)
Normal copy to between src and dst.
Definition: Promotion.cpp:503
void populateDecomposeConvolutionPatterns(RewritePatternSet &patterns, PatternBenefit benefit=1)
Linalg decompose convolutions patterns.
void populateConvolutionVectorizationPatterns(RewritePatternSet &patterns, PatternBenefit benefit=1)
Populate patterns for vectorizing low-D convolution ops.
LogicalResult vectorizeCopy(RewriterBase &builder, memref::CopyOp copyOp)
Emit a suitable vector form for a Copy op with fully static shape.
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:50
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:111
FailureOr< LinalgLoops > linalgOpToAffineLoops(RewriterBase &rewriter, LinalgOp linalgOp)
Emit a loop nest of affine.for with the proper body for linalgOp.
Definition: Loops.cpp:363
void populateEraseUnusedOperandsAndResultsPatterns(RewritePatternSet &patterns)
Pattern to remove dead operands and results of linalg.generic operations.
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:352
FailureOr< SplitReductionResult > splitReduction(RewriterBase &b, LinalgOp op, const ControlSplitReductionFn &controlSplitReductionFn, bool useAlloc=false)
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:265
void populateDecomposeLinalgOpsPattern(RewritePatternSet &patterns, bool removeDeadArgsAndResults=true)
Populate patterns for splitting a LinalgOp with multiple statements within its payload into multiple ...
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:647
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:769
LogicalResult dropUnitDims(RewriterBase &rewriter, GenericOp genericOp, const ControlDropUnitDims &options)
FailureOr< PackResult > pack(RewriterBase &rewriter, linalg::LinalgOp linalgOp, ArrayRef< OpFoldResult > packedSizes)
Implement packing of a single LinalgOp by packedSizes.
Definition: Transforms.cpp:480
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:872
std::function< SmallVector< ReassociationIndices >(linalg::LinalgOp)> GetCollapsableDimensionsFn
Function type to control generic op dimension collapsing.
Definition: Transforms.h:1631
FailureOr< LowerPackResult > lowerPack(RewriterBase &rewriter, tensor::PackOp packOp)
Rewrite pack as pad + reshape + transpose.
Definition: Transforms.cpp:219
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:421
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:443
LogicalResult deallocateWorkgroupMemory(OpBuilder &, Value)
In case of GPU group memory there is no need to deallocate.
Definition: Promotion.cpp:479
FailureOr< Operation * > transposeMatmul(RewriterBase &rewriter, linalg::MatmulOp op, bool transposeLHS=true)
Convert Linalg matmul ops to transposed variants.
void populateLinalgNamedOpsGeneralizationPatterns(RewritePatternSet &patterns)
Linalg generalization patterns.
std::function< std::optional< BlockPackMatmulOptions >(linalg::LinalgOp)> ControlBlockPackMatmulFn
Function type which is used to control matmul packing.
Definition: Transforms.h:1198
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:59
RewritePatternSet getLinalgTilingCanonicalizationPatterns(MLIRContext *ctx)
Canonicalization patterns relevant to apply after tiling patterns.
Definition: Tiling.cpp:906
FailureOr< CollapseResult > collapseOpIterationDims(LinalgOp op, ArrayRef< ReassociationIndices > foldedIterationDims, RewriterBase &rewriter)
Collapses dimensions of linalg.generic/linalg.copy operation.
FailureOr< PackTransposeResult > packTranspose(RewriterBase &rewriter, tensor::PackOp packOp, linalg::LinalgOp linalgOp, tensor::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:678
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:137
FailureOr< LinalgLoops > linalgOpToParallelLoops(RewriterBase &rewriter, LinalgOp linalgOp)
Emit a loop nest of scf.parallel with the proper body for linalgOp.
Definition: Loops.cpp:375
Include the generated interface declarations.
ArrayRef< int64_t > ReassociationIndicesRef
This class represents an efficient way to signal success or failure.
Definition: LogicalResult.h:26
OpInterfaceRewritePattern is a wrapper around RewritePattern that allows for matching and rewriting a...
Definition: PatternMatch.h:373
OpRewritePattern is a wrapper around RewritePattern that allows for matching and rewriting against an...
Definition: PatternMatch.h:358
SmallVector< int64_t, 3 > mnkOrder
Permutation of matmul (M, N, K) dimensions order.
Definition: Transforms.h:1178
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:1168
bool rhsTransposeOuterBlocks
Transpose RHS outer block layout [KB][NB] -> [NB][KB].
Definition: Transforms.h:1187
bool lhsTransposeInnerBlocks
Transpose LHS inner block layout [mb][kb] -> [kb][mb].
Definition: Transforms.h:1184
SmallVector< int64_t, 3 > mnkPaddedSizesNextMultipleOf
Next multiples of the packing sizes.
Definition: Transforms.h:1175
bool lhsTransposeOuterBlocks
Transpose LHS outer block layout [MB][KB] -> [KB][MB].
Definition: Transforms.h:1181
bool allowPadding
If true, allows packing of dimensions that only partially fit into the block factors.
Definition: Transforms.h:1172
bool rhsTransposeInnerBlocks
Transpose RHS inner block layout [kb][nb] -> [nb][kb].
Definition: Transforms.h:1190
bool bufferizeDestinationOnly
If set to "true", only the destination tensor operands are bufferized to a new allocation (and wrappe...
Definition: Transforms.h:65
bool emitDealloc
If set to "true", a memref.dealloc operation will be emitted for each allocated buffer.
Definition: Transforms.h:71
SmallVector< Value > results
Definition: Transforms.h:1083
Transformation to drop unit-extent dimensions from linalg.generic operations.
Definition: Transforms.h:474
RankReductionStrategy rankReductionStrategy
Definition: Transforms.h:477
std::function< SmallVector< unsigned >(Operation *)> ControlFnTy
Definition: Transforms.h:480
Vectorization pattern for memref::CopyOp.
Definition: Transforms.h:1399
LogicalResult matchAndRewrite(memref::CopyOp copyOp, PatternRewriter &rewriter) const override
Definition: Transforms.cpp:920
LogicalResult matchAndRewrite(Conv2DOp convOp, PatternRewriter &rewriter) const override
Definition: Transforms.h:1367
FailureOr< Conv1DOp > returningMatchAndRewrite(Conv2DOp convOp, PatternRewriter &rewriter) const
DownscaleConv2DOp(MLIRContext *context, PatternBenefit benefit=1)
Definition: Transforms.h:1361
Rewrites 2-D depthwise convolution ops with size-1 (w, kw) or (h, kh) dimensions into 1-D depthwise c...
Definition: Transforms.h:1345
FailureOr< DepthwiseConv1DNwcWcOp > returningMatchAndRewrite(DepthwiseConv2DNhwcHwcOp convOp, PatternRewriter &rewriter) const
LogicalResult matchAndRewrite(DepthwiseConv2DNhwcHwcOp convOp, PatternRewriter &rewriter) const override
Definition: Transforms.h:1354
DownscaleDepthwiseConv2DNhwcHwcOp(MLIRContext *context, PatternBenefit benefit=1)
Definition: Transforms.h:1346
Rewrites 2-D convolution ops with size-1 window dimensions into 1-D convolution ops.
Definition: Transforms.h:1325
LogicalResult matchAndRewrite(Conv2DOp convOp, PatternRewriter &rewriter) const override
Definition: Transforms.h:1331
FailureOr< Conv1DOp > returningMatchAndRewrite(Conv2DOp convOp, PatternRewriter &rewriter) const
Fuse two linalg.generic operations that have a producer-consumer relationship captured through fusedO...
Definition: Transforms.h:498
llvm::DenseMap< Value, Value > replacements
Definition: Transforms.h:500
static llvm::SmallDenseSet< int > getPreservedProducerResults(GenericOp producer, GenericOp consumer)
Returns a set of indices of the producer's results which would be preserved after the fusion.
Rewrite extract_slice(tensor.pad(x)) into tensor.pad(extract_slice(x)).
Definition: Transforms.h:1509
std::function< std::optional< bool >(tensor::ExtractSliceOp)> ControlFn
A function to control pattern application and rewrite logic.
Definition: Transforms.h:1519
LogicalResult matchAndRewrite(tensor::ExtractSliceOp sliceOp, PatternRewriter &rewriter) const override
Definition: Transforms.cpp:999
ExtractSliceOfPadTensorSwapPattern(MLIRContext *context, ControlFn controlFn=nullptr, PatternBenefit benefit=1)
Definition: Transforms.h:1521
Transformation information returned after reduction tiling.
Definition: Transforms.h:874
Operation * parallelTiledOp
The partial reduction tiled op generated.
Definition: Transforms.h:876
SmallVector< Value > initialValues
Initial values used for partial reductions.
Definition: Transforms.h:880
scf::ForallOp loops
The scf.forall operation that iterate over the tiles.
Definition: Transforms.h:882
Operation * mergeOp
The final reduction operation merging all the partial reductions.
Definition: Transforms.h:878
Rewrite a TilingInterface op to a tiled scf.forall, applying tiling by numThreads.
Definition: Transforms.h:857
Rewrites a tensor::PackOp into a sequence of tensor.pad + linalg.transpose + tensor....
Definition: Transforms.h:1432
LogicalResult matchAndRewrite(tensor::PackOp packOp, PatternRewriter &rewriter) const override
Rewrites a tensor::UnPackOp into a sequence of rank-reduced extract_slice op.
Definition: Transforms.h:1442
LogicalResult matchAndRewrite(tensor::UnPackOp unpackOp, PatternRewriter &rewriter) const override
Rewrite a tensor::PadOp into a sequence of EmptyOp, FillOp and InsertSliceOp.
Definition: Transforms.h:1412
LogicalResult matchAndRewrite(tensor::PadOp padOp, PatternRewriter &rewriter) const override
Definition: Transforms.cpp:944
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:927
GeneralizePadOpPattern(MLIRContext *context, OptimizeCopyFn optimizeCopyFn=nullptr, PatternBenefit benefit=1)
Definition: Transforms.h:1413
Match and rewrite for the pattern:
Definition: Transforms.h:1472
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:1500
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:1382
LogicalResult matchAndRewrite(LinalgOp op, PatternRewriter &rewriter) const override
Definition: Transforms.h:1392
FailureOr< GenericOp > returningMatchAndRewrite(LinalgOp op, PatternRewriter &rewriter) const
matchAndRewrite implementation that returns the significant transformed pieces of IR.
Definition: Transforms.h:1388
Options that allow distribution of loops generated in Linalg transforms to processors while generatin...
Definition: Utils.h:305
SmallVector< Attribute > paddingValues
A padding value for every operand.
Definition: Transforms.h:281
LinalgPaddingOptions & setPadToMultipleOf(ArrayRef< int64_t > m)
Definition: Transforms.h:294
SmallVector< bool > packPaddings
A flag for every operand to mark the PadOp as nofold which enables packing for statically shaped oper...
Definition: Transforms.h:300
std::optional< SmallVector< int64_t > > padToMultipleOf
A list of multiples to which each padding dimension should be padded to.
Definition: Transforms.h:293
LinalgPaddingOptions & setPaddingDimensions(ArrayRef< int64_t > pd)
Definition: Transforms.h:288
LinalgPaddingOptions & setTransposePaddings(ArrayRef< SmallVector< int64_t >> tp)
Definition: Transforms.h:315
SmallVector< SmallVector< int64_t > > transposePaddings
A permutation vector for every operand used to transpose the packed PadOp results.
Definition: Transforms.h:313
LinalgPaddingOptions & setPaddingValues(ArrayRef< Attribute > pv)
Definition: Transforms.h:282
LinalgPaddingOptions & setPackPaddings(ArrayRef< bool > pp)
Definition: Transforms.h:301
LinalgPaddingOptions & setCopyBackOp(CopyBackOp op)
Definition: Transforms.h:327
LinalgPaddingOptions & setHoistPaddings(ArrayRef< int64_t > hp)
Definition: Transforms.h:307
SmallVector< int64_t > hoistPaddings
A number of loops to hoist the PadOp out for every operand.
Definition: Transforms.h:306
SmallVector< int64_t > paddingDimensions
A list of iterator dimensions to pad.
Definition: Transforms.h:287
CopyBackOp copyBackOp
The op to be used for copying the padded result to the original destination tensor.
Definition: Transforms.h:326
std::optional< unsigned > alignment
Alignment of promoted buffer. If std::nullopt do not specify alignment.
Definition: Transforms.h:385
LinalgPromotionOptions & setUseFullTileBuffersByDefault(bool use)
Definition: Transforms.h:380
bool useAlloca
Use alloca with the default allocation scheme.
Definition: Transforms.h:398
LinalgPromotionOptions & setAlignment(unsigned align)
Definition: Transforms.h:386
std::optional< Attribute > memorySpace
Memory space of promoted buffer.
Definition: Transforms.h:392
std::optional< CopyCallbackFn > copyOutFn
Definition: Transforms.h:418
std::optional< CopyCallbackFn > copyInFn
Callback function to do the copy of data to and from the promoted subview.
Definition: Transforms.h:417
LinalgPromotionOptions & setUseAlloca(bool use)
Definition: Transforms.h:399
std::optional< DenseSet< unsigned > > operandsToPromote
Indices of subViews to promote.
Definition: Transforms.h:357
LinalgPromotionOptions & setCopyInOutFns(CopyCallbackFn const &copyIn, CopyCallbackFn const &copyOut)
Definition: Transforms.h:419
LinalgPromotionOptions & setUseFullTileBuffers(ArrayRef< bool > useFullTiles)
Definition: Transforms.h:369
std::optional< AllocBufferCallbackFn > allocationFn
Callback function to do the allocation of the promoted buffer.
Definition: Transforms.h:406
bool useFullTileBuffersDefault
If true all operands unspecified by useFullTileBuffers will use the full view, otherwise the partial ...
Definition: Transforms.h:379
std::optional< DeallocBufferCallbackFn > deallocationFn
Definition: Transforms.h:407
LinalgPromotionOptions & setMemorySpace(Attribute memorySpc)
Definition: Transforms.h:393
LinalgPromotionOptions & setAllocationDeallocationFns(AllocBufferCallbackFn const &allocFn, DeallocBufferCallbackFn const &deallocFn)
Definition: Transforms.h:409
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:368
LinalgPromotionOptions & setOperandsToPromote(ArrayRef< int64_t > operands)
Definition: Transforms.h:358
std::optional< LinalgLoopDistributionOptions > tileDistribution
When specified, specifies distribution of generated tile loops to processors.
Definition: Transforms.h:271
LinalgTilingAndFusionOptions & setTileSizes(ArrayRef< int64_t > ts)
Definition: Transforms.h:263
SmallVector< int64_t > tileInterchange
Tile interchange used to permute the tile loops.
Definition: Transforms.h:268
LinalgTilingAndFusionOptions & setDistributionOptions(LinalgLoopDistributionOptions distributionOptions)
Definition: Transforms.h:273
SmallVector< int64_t > tileSizes
Tile sizes used to tile the root operation.
Definition: Transforms.h:262
LinalgTilingOptions & setLoopType(LinalgTilingLoopType lt)
Definition: Transforms.h:227
LinalgTilingOptions & setDistributionTypes(ArrayRef< StringRef > types)
Definition: Transforms.h:245
LinalgTilingOptions & setInterchange(ArrayRef< unsigned > interchange)
Definition: Transforms.h:219
LinalgTilingLoopType loopType
The type of tile loops to generate.
Definition: Transforms.h:225
LinalgTilingOptions & setTileSizeComputationFunction(TileSizeComputationFunction fun)
Definition: Transforms.h:196
LinalgTilingOptions & setTileSizes(const SmallVector< Value, 4 > &ts)
Set the tileSizeComputationFunction to return the values ts.
Definition: Transforms.h:203
LinalgTilingOptions & setPeeledLoops(ArrayRef< int64_t > loops)
Definition: Transforms.h:253
SmallVector< int64_t > peeledLoops
Peel the specified loops.
Definition: Transforms.h:251
LinalgTilingOptions & setDistributionOptions(LinalgLoopDistributionOptions distributionOptions)
Definition: Transforms.h:237
SmallVector< unsigned, 4 > interchangeVector
The interchange vector to reorder the tiled loops.
Definition: Transforms.h:217
TileSizeComputationFunction tileSizeComputationFunction
Computation function that returns the tile sizes for each operation.
Definition: Transforms.h:193
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:234
SmallVector< StringRef, 2 > distributionTypes
Specification markers of how to distribute the linalg.tiled_loop.
Definition: Transforms.h:243
linalg::TransposeOp transposeOp
Definition: Transforms.h:1102
tensor::ExpandShapeOp expandShapeOp
Definition: Transforms.h:1101
tensor::ExtractSliceOp extractSliceOp
Definition: Transforms.h:1113
linalg::TransposeOp transposeOp
Definition: Transforms.h:1111
tensor::CollapseShapeOp collapseShapeOp
Definition: Transforms.h:1112
A description of a multi-size tiling comprising tile sizes and numbers of tiles, expressed as Values ...
Definition: Transforms.h:810
Struct to hold the result of a pack call.
Definition: Transforms.h:1121
linalg::LinalgOp packedLinalgOp
Definition: Transforms.h:1123
SmallVector< tensor::PackOp > packOps
Definition: Transforms.h:1122
SmallVector< tensor::UnPackOp > unPackOps
Definition: Transforms.h:1124
Struct to hold the result of a packTranspose call.
Definition: Transforms.h:1133
linalg::LinalgOp transposedLinalgOp
Definition: Transforms.h:1135
tensor::UnPackOp transposedUnPackOp
Definition: Transforms.h:1136
Create a new buffer using the allocationFn provided.
Definition: Transforms.h:707
Split Reduction options.
Definition: Transforms.h:428
Apply transformation to split the single linalg op reduction into a parallel and reduction dimension.
Definition: Transforms.h:1011
Perform standalone tiling of a single LinalgOp by tileSizes.
Definition: Transforms.h:668
SmallVector< Operation *, 8 > loops
Definition: Transforms.h:670
SmallVector< Value, 4 > tensorResults
Definition: Transforms.h:671
T lowTripCount
Number of tiles associated with each size.
Definition: Transforms.h:802
Helper struct to hold the results of building a packing loop nest.
Definition: Transforms.h:538
SmallVector< OpFoldResult > strides
Definition: Transforms.h:539
SmallVector< Value > leadingPackedTensorIndexings
Definition: Transforms.h:540
SmallVector< Value > clonedLoopIvs
Definition: Transforms.h:540
SmallVector< OpFoldResult > sizes
Definition: Transforms.h:539
SmallVector< OpFoldResult > offsets
Definition: Transforms.h:539