MLIR  19.0.0git
Utils.h
Go to the documentation of this file.
1 //===- Utils.h - Utilities to support the Linalg dialect --------*- 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_UTILS_UTILS_H
10 #define MLIR_DIALECT_LINALG_UTILS_UTILS_H
11 
15 #include "llvm/ADT/StringSet.h"
16 #include <optional>
17 
18 namespace mlir {
19 class AffineExpr;
20 class AffineMap;
21 class PatternRewriter;
22 
23 namespace affine {
24 class AffineForOp;
25 } // namespace affine
26 
27 namespace tensor {
28 class ExtractSliceOp;
29 } // namespace tensor
30 
31 namespace linalg {
32 
33 //===----------------------------------------------------------------------===//
34 // Utilities for inferring various semantics properties of Linalg ops.
35 //===----------------------------------------------------------------------===//
36 
37 //===----------------------------------------------------------------------===//
38 // General utilities
39 //===----------------------------------------------------------------------===//
40 
41 /// Check if all indexing maps are projected permutations.
42 bool allIndexingsAreProjectedPermutation(LinalgOp op);
43 
44 /// Detect whether `r` has only ConstantOp, ElementwiseMappable and YieldOp.
45 bool hasOnlyScalarElementwiseOp(Region &r);
46 
47 /// Check if a LinalgOp is an element-wise operation.
48 bool isElementwise(LinalgOp op);
49 
50 /// Check if iterator type has "parallel" semantics.
51 bool isParallelIterator(utils::IteratorType iteratorType);
52 
53 /// Check if iterator type has "reduction" semantics.
54 bool isReductionIterator(utils::IteratorType iteratorType);
55 
56 /// Create a tensor::PadOp that pads `source` to the size of the statically
57 /// sized `type` whose static sizes are assumed to be greater than the dynamic
58 /// `source` size. The padding introduces trailing `pad` values until the
59 /// target size is met. If `source` is defined by one or more LinalgOps that
60 /// have been padded with the same value and sizes, return their padded result
61 /// instead of creating a tensor::PadOp.
62 ///
63 /// Example:
64 /// ```
65 /// %0 = tensor.extract_slice %arg0 [%iv0, %iv1] [%sz0, %sz1]
66 /// %1 = tensor.pad %0 low[0, 0] high[...] { tensor.yield %cst }
67 /// %2 = linalg.matmul ins(...) outs(%1)
68 /// %3 = tensor.extract_slice %2 [0, 0] [%sz0, %sz1]
69 /// ```
70 /// makeComposedPadHighOp(source=%3, pad=%cst) returns %2
71 /// makeComposedPadHighOp(source=%3, pad=%other_cst) returns %4
72 /// ```
73 /// %4 = tensor.pad %3 low[0, 0] high[...] { tensor.yield %other_cst }
74 /// ```
75 Value makeComposedPadHighOp(OpBuilder &b, Location loc, RankedTensorType type,
76  Value source, Value pad, bool nofold);
77 
78 /// Returns a GenericOp that transposes `inputTensor` into `outputTensor`
79 /// using `transposeVector` to permute the `inputTensor` dimensions.
80 GenericOp makeTransposeOp(OpBuilder &b, Location loc, Value inputTensor,
81  Value outputTensor,
82  ArrayRef<int64_t> transposeVector);
83 
84 /// Returns GenericOp that copies an n-D memref. Unlike the current
85 /// implementation of memref::CopyOp, this op can further tile, lower to loops
86 /// or vectorize.
87 GenericOp makeMemRefCopyOp(OpBuilder &b, Location loc, Value from, Value to);
88 
89 /// Get the reassociation maps to fold the result of a extract_slice (or
90 /// source of a insert_slice) operation with given offsets, and sizes to its
91 /// rank-reduced version. This is only done for the cases where the size is 1
92 /// and offset is 0. Strictly speaking the offset 0 is not required in
93 /// general, but non-zero offsets are not handled by SPIR-V backend at this
94 /// point (and potentially cannot be handled).
95 std::optional<SmallVector<ReassociationIndices>>
96 getReassociationMapForFoldingUnitDims(ArrayRef<OpFoldResult> mixedSizes);
97 
98 //===----------------------------------------------------------------------===//
99 // Fusion / Tiling utilities
100 //===----------------------------------------------------------------------===//
101 
102 /// The type of loops to be generated during tiling.
104  Loops = 0,
105  AffineLoops = 1,
106  ParallelLoops = 2
107 };
108 
109 /// Computes tile offsets, given a list of loop `ivs` and `tileSizes`. In case
110 /// a tile size is zero (i.e., no tiling), the corresponding offset is also
111 /// zero.
114  ArrayRef<OpFoldResult> tileSizes);
115 
116 /// Computes tile sizes, given a list of `tileSizes` and dimension
117 /// sizes (`sizeBounds`). In case a tile size is zero (i.e., no tiling), the
118 /// corresponding result size is the corresponding value from `sizeBounds`.
119 /// Note: The returned tile sizes are closed intervals.
121  ArrayRef<OpFoldResult> tileSizes,
122  ArrayRef<OpFoldResult> sizeBounds);
123 
124 /// Returns the list of tensor output types produced when the given structured
125 /// operation `op` is applied to the given `operands`. Note that `operands`
126 /// are not necessarily the actual operands of `op`.
127 SmallVector<Type> getTensorOutputTypes(LinalgOp op, ValueRange operands);
128 
129 /// Creates `insert_slice` ops that insert `results` back into larger tensors
130 /// they were originally extracted from with `extract_slice` before being
131 /// passed as `operands` to the given structured operation `op` or its clone.
132 /// Note that `operands` are not necessarily the actual operands of `op`, the
133 /// operation serves only as metadata container for operand types and
134 /// positions.
136  LinalgOp op, ValueRange operands,
137  ValueRange results);
138 
139 /// A struct containg offsets-sizes-strides arguments of the tiled shape.
144 };
145 
146 /// Computes SliceParameters for a single `valueToTile` assuming that its user
147 /// is being tiled with the given loop bounds `lbs` and `ubs` and the tile
148 /// sizes `tileSizes`.
149 ///
150 /// `omitPartialTileCheck` controls whether to omit the partial/boundary tile
151 /// condition check in cases where we statically know that it is unnecessary.
153 computeSliceParameters(OpBuilder &builder, Location loc, Value valueToTile,
154  ArrayRef<OpFoldResult> tileSizes, AffineMap map,
156  ArrayRef<OpFoldResult> subShapeSizes,
157  bool omitPartialTileCheck);
158 
159 /// Computes SliceParamaters for all `valuesToTile` of the given `linalgOp`,
160 /// assuming `linalgOp` is being fused into a loop nest. Calls
161 /// `computeSliceParameters` for every individual value.
162 ///
163 /// Note that a constant zero in `tileSizes` means no tiling at that implicit
164 /// loop. The number of non-zero values in `tileSizes` should be equal to the
165 /// number of values in `ivs`.
166 ///
167 /// Some of the `valuesToTile` won't be affected by tiling. For these values,
168 /// std::nullopt will be returned.
170 computeAllSliceParameters(OpBuilder &builder, Location loc, LinalgOp linalgOp,
171  ValueRange valuesToTile, ArrayRef<OpFoldResult> ivs,
172  ArrayRef<OpFoldResult> tileSizes,
173  ArrayRef<OpFoldResult> sizeBounds,
174  bool omitPartialTileCheck);
175 
176 /// Creates an extract_slice/subview op for a single `valueToTile` with
177 /// `builder`. This new operation extracts a tile of `valueToTile`, starting
178 /// at offsets `lbs` and with sizes `subShapeSizes`. `omitPartialTileCheck`
179 /// controls whether to omit the partial/boundary tile condition check in
180 /// cases where we statically know that it is unnecessary.
181 Value makeTiledShape(OpBuilder &builder, Location loc, Value valueToTile,
182  ArrayRef<OpFoldResult> tileSizes, AffineMap map,
184  ArrayRef<OpFoldResult> subShapeSizes,
185  bool omitPartialTileCheck);
186 
187 /// Creates extract_slice/subview ops for all `valuesToTile` of the given
188 /// `linalgOp` with `builder`, assuming `linalgOp` is being fused into a loop
189 /// nest for tiling with the given induction variables `ivs` and tile sizes
190 /// `tileSizes`. `sizeBounds` are the iteration space bounds for *all* the
191 /// implicit loops in `linalgOp`. `omitPartialTileCheck` controls whether to
192 /// omit the partial/boundary tile condition check in cases where we
193 /// statically know that it is unnecessary.
194 ///
195 /// Note that a constant zero in `tileSizes` means no tiling at that implicit
196 /// loop. The number of non-zero values in `tileSizes` should be equal to the
197 /// number of values in `ivs`.
199  LinalgOp linalgOp, ValueRange valuesToTile,
201  ArrayRef<OpFoldResult> tileSizes,
202  ArrayRef<OpFoldResult> sizeBounds,
203  bool omitPartialTileCheck);
204 
205 /// Add the specified offsets to any `linalg.index` ops contained in the given
206 /// `linalgOp`. The offsets are provided in the same order as iteration space
207 /// dimensions. Null offests are assumed to be zero.
208 void offsetIndices(OpBuilder &b, LinalgOp linalgOp,
209  ArrayRef<OpFoldResult> offests);
210 void offsetIndices(RewriterBase &b, LinalgOp linalgOp,
211  ArrayRef<OpFoldResult> offests);
212 
213 /// A struct containing the Linalg producer before and after fusion.
214 /// When operating on tensors, `fusedProducer` may feed into a `tensor.cast`
215 /// op before the consumer Linalg op, until enough canonicalizations have
216 /// applied.
217 struct FusionInfo {
219  LinalgOp fusedProducer;
220 };
221 
222 /// This implements the fusion part of the "tileAndFuse on tensors"
223 /// transformation and thus requires the `consumerOpOperand` to be a
224 /// `extract_slice` op (generally obtained by applying the tiling
225 /// transformation).
226 FailureOr<FusionInfo> fuseProducerOfTensor(OpBuilder &b,
227  OpOperand &consumerOpOperand);
228 
229 /// This implements the fusion part of the "tileAndFuse on tensors"
230 /// transformation and thus requires the `consumerOpOperand` to be a
231 /// `extract_slice` op (generally obtained by applying the tiling
232 /// transformation). Assumes `producerOfTensor` is a Linalg op that produces
233 /// `consumerOpOperand`.
234 FailureOr<FusionInfo> fuseProducerOfTensor(OpBuilder &b,
235  OpResult producerOpResult,
236  OpOperand &consumerOpOperand);
237 
238 //===----------------------------------------------------------------------===//
239 // Distribution utilities
240 //===----------------------------------------------------------------------===//
241 
242 /// Scheme used to distribute loops to processors.
243 enum class DistributionMethod {
244  /// Cyclic distribution where no assumption is made about the dynamic
245  /// relationship between number of processors and number of iterations of
246  /// the
247  /// distributed loop. Distributes the following loop
248  ///
249  /// scf.parallel (%iv) = (%lb) to (%ub) step (%step)
250  ///
251  /// to
252  ///
253  /// scf.parallel(%iv)= (%lb + %procId * %step) to (%ub) step (%step *
254  /// %nprocs)
255  Cyclic = 0,
256 
257  /// Cyclic distribution where the number of processors can be assumed to be
258  /// more than or equal to the number of iterations of the distributed loop.
259  /// In
260  /// such cases, a simple in-bounds check is enough (instead of materializing
261  /// a
262  /// loop). Distributes the following loop
263  ///
264  /// scf.parallel (%iv) = (%lb) to (%ub) step (%step)
265  ///
266  /// to
267  ///
268  /// %iv = %lb + %procId * %step
269  /// %cond = arith.cmpi "slt", %iv, %ub
270  /// scf.if %cond {
271  /// ...
272  /// }
274 
275  /// Cyclic distribution where the number of processors can be assumed to be
276  /// equal to the number of iterations of the distributed loop. In such
277  /// cases,
278  /// no bounds check is needed. Distributes the following loop
279  ///
280  /// scf.parallel (%iv) = (%lb) to (%ub) step (%step)
281  ///
282  /// to
283  ///
284  /// %iv = %lb + %procId * %step
286 
287  /// No Distribution.
288  None = 3
289 };
290 
291 /// Callback function type used to get processor ID, and number of processors
292 /// used for distribution for all parallel loops generated.
293 struct ProcInfo {
297 };
298 using ProcInfoCallBackFn = std::function<SmallVector<ProcInfo>(
299  OpBuilder &b, Location loc, ArrayRef<Range> parallelLoopRanges)>;
300 
301 /// Options that allow distribution of loops generated in Linalg transforms to
302 /// processors while generating the loops.
304  /// Callback function that returns the Values for processor ID (`procId`),
305  /// and number of processors (`nprocs`) used to execute the parallel loops.
306  /// The number of `{procId, nprocs}` pairs returned must be equal to the
307  /// number of `parallelLoopRanges` passed into the callback. The
308  /// `parallelLoopRanges` are ranges of the outer parallel loops of the
309  /// operation that do have non-zero tile sizes specified.
311 };
312 
313 /// Update the `lb`, `ub` and `step` to get per processor `lb`, `ub` and
314 /// `step`.
316  Value procId, Value nprocs, Value &lb,
317  Value &ub, Value &step);
318 
319 //===----------------------------------------------------------------------===//
320 // Fusion on tensor utilities
321 //===----------------------------------------------------------------------===//
322 
323 //===----------------------------------------------------------------------===//
324 // Generic op region utilities
325 //===----------------------------------------------------------------------===//
326 
327 /// A struct containing common matchers over linalg op's region.
329  enum class BinaryOpKind {
330  IAdd,
331  };
332 
333  /// Matches the given linalg op if its body is performing binary operation
334  /// on int or float scalar values and returns the binary op kind.
335  ///
336  /// The linalg op's region is expected to be
337  /// ```
338  /// {
339  /// ^bb(%a: <scalar-type>, %b: <scalar-type>):
340  /// %0 = <binary-op> %a, %b: <scalar-type>
341  /// linalg.yield %0: <scalar-type>
342  /// }
343  /// ```
344  static std::optional<BinaryOpKind> matchAsScalarBinaryOp(GenericOp op);
345 };
346 
347 //===----------------------------------------------------------------------===//
348 // Loop nest utilities
349 //===----------------------------------------------------------------------===//
350 
351 /// Utility class used to generate nested loops with ranges described by
352 /// `loopRanges` and loop type described by the `iteratorTypes`.
353 /// `bodyBuilderFn` is used to generate the body of the innermost loop. It is
354 /// passed a range of loop induction variables and a range of operand values
355 /// to use.
356 template <typename LoopTy>
358  static void doit(OpBuilder &b, Location loc, ArrayRef<Range> loopRanges,
359  LinalgOp linalgOp,
360  ArrayRef<utils::IteratorType> iteratorTypes,
363  bodyBuilderFn,
364  ArrayRef<linalg::ProcInfo> procInfo = {});
365 };
366 
367 /// Returns an attribute list that excludes pre-defined attributes.
368 template <typename OpTy>
370  auto elidedAttrs = llvm::to_vector(op.getAttributeNames());
371  if (isa<linalg::LinalgOp>(op.getOperation()))
372  elidedAttrs.push_back(LinalgDialect::kMemoizedIndexingMapsAttrName);
373  return getPrunedAttributeList(op, elidedAttrs);
374 }
375 
376 } // namespace linalg
377 } // namespace mlir
378 
379 #endif // MLIR_DIALECT_LINALG_UTILS_UTILS_H
A multi-dimensional affine map Affine map's are immutable like Type's, and they are uniqued.
Definition: AffineMap.h:46
This class defines the main interface for locations in MLIR and acts as a non-nullable wrapper around...
Definition: Location.h:63
This class helps build Operations.
Definition: Builders.h:209
This class represents an operand of an operation.
Definition: Value.h:267
This is a value defined by a result of an operation.
Definition: Value.h:457
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 provides an abstraction over the different types of ranges over Values.
Definition: ValueRange.h:381
This class represents an instance of an SSA value in the MLIR system, representing a computable value...
Definition: Value.h:96
FailureOr< FusionInfo > fuseProducerOfTensor(OpBuilder &b, OpOperand &consumerOpOperand)
This implements the fusion part of the "tileAndFuse on tensors" transformation and thus requires the ...
Definition: Fusion.cpp:227
SmallVector< Value > makeTiledShapes(OpBuilder &builder, Location loc, LinalgOp linalgOp, ValueRange valuesToTile, ArrayRef< OpFoldResult > ivs, ArrayRef< OpFoldResult > tileSizes, ArrayRef< OpFoldResult > sizeBounds, bool omitPartialTileCheck)
Creates extract_slice/subview ops for all valuesToTile of the given linalgOp with builder,...
Definition: Utils.cpp:829
bool allIndexingsAreProjectedPermutation(LinalgOp op)
Check if all indexing maps are projected permutations.
Definition: Utils.cpp:149
bool isParallelIterator(utils::IteratorType iteratorType)
Check if iterator type has "parallel" semantics.
Definition: Utils.cpp:184
SmallVector< OpFoldResult > computeTileSizes(OpBuilder &b, Location loc, ArrayRef< OpFoldResult > tileSizes, ArrayRef< OpFoldResult > sizeBounds)
Computes tile sizes, given a list of tileSizes and dimension sizes (sizeBounds).
Definition: Utils.cpp:724
GenericOp makeMemRefCopyOp(OpBuilder &b, Location loc, Value from, Value to)
Returns GenericOp that copies an n-D memref.
Definition: Utils.cpp:288
SmallVector< OpFoldResult > computeTileOffsets(OpBuilder &b, Location loc, ArrayRef< OpFoldResult > ivs, ArrayRef< OpFoldResult > tileSizes)
Computes tile offsets, given a list of loop ivs and tileSizes.
Definition: Utils.cpp:710
LinalgTilingLoopType
The type of loops to be generated during tiling.
Definition: Utils.h:103
bool isReductionIterator(utils::IteratorType iteratorType)
Check if iterator type has "reduction" semantics.
Definition: Utils.cpp:188
bool hasOnlyScalarElementwiseOp(Region &r)
Detect whether r has only ConstantOp, ElementwiseMappable and YieldOp.
Definition: Utils.cpp:155
GenericOp makeTransposeOp(OpBuilder &b, Location loc, Value inputTensor, Value outputTensor, ArrayRef< int64_t > transposeVector)
Returns a GenericOp that transposes inputTensor into outputTensor using transposeVector to permute th...
Definition: Utils.cpp:252
std::function< SmallVector< ProcInfo >(OpBuilder &b, Location loc, ArrayRef< Range > parallelLoopRanges)> ProcInfoCallBackFn
Definition: Utils.h:299
SmallVector< NamedAttribute > getPrunedAttributeList(OpTy op)
Returns an attribute list that excludes pre-defined attributes.
Definition: Utils.h:369
std::optional< SmallVector< ReassociationIndices > > getReassociationMapForFoldingUnitDims(ArrayRef< OpFoldResult > mixedSizes)
Get the reassociation maps to fold the result of a extract_slice (or source of a insert_slice) operat...
Definition: Utils.cpp:886
DistributionMethod
Scheme used to distribute loops to processors.
Definition: Utils.h:243
@ CyclicNumProcsGeNumIters
Cyclic distribution where the number of processors can be assumed to be more than or equal to the num...
@ Cyclic
Cyclic distribution where no assumption is made about the dynamic relationship between number of proc...
@ CyclicNumProcsEqNumIters
Cyclic distribution where the number of processors can be assumed to be equal to the number of iterat...
SmallVector< Value > insertSlicesBack(OpBuilder &builder, Location loc, LinalgOp op, ValueRange operands, ValueRange results)
Creates insert_slice ops that insert results back into larger tensors they were originally extracted ...
Definition: Utils.cpp:749
bool isElementwise(LinalgOp op)
Check if a LinalgOp is an element-wise operation.
Definition: Utils.cpp:169
void offsetIndices(OpBuilder &b, LinalgOp linalgOp, ArrayRef< OpFoldResult > offests)
Add the specified offsets to any linalg.index ops contained in the given linalgOp.
Definition: Utils.cpp:850
SmallVector< std::optional< SliceParameters > > computeAllSliceParameters(OpBuilder &builder, Location loc, LinalgOp linalgOp, ValueRange valuesToTile, ArrayRef< OpFoldResult > ivs, ArrayRef< OpFoldResult > tileSizes, ArrayRef< OpFoldResult > sizeBounds, bool omitPartialTileCheck)
Computes SliceParamaters for all valuesToTile of the given linalgOp, assuming linalgOp is being fused...
Definition: Utils.cpp:778
Value makeComposedPadHighOp(OpBuilder &b, Location loc, RankedTensorType type, Value source, Value pad, bool nofold)
Create a tensor::PadOp that pads source to the size of the statically sized type whose static sizes a...
Definition: Utils.cpp:192
void updateBoundsForCyclicDistribution(OpBuilder &builder, Location loc, Value procId, Value nprocs, Value &lb, Value &ub, Value &step)
Update the lb, ub and step to get per processor lb, ub and step.
Definition: Utils.cpp:387
SmallVector< Type > getTensorOutputTypes(LinalgOp op, ValueRange operands)
Returns the list of tensor output types produced when the given structured operation op is applied to...
Definition: Utils.cpp:740
Value makeTiledShape(OpBuilder &builder, Location loc, Value valueToTile, ArrayRef< OpFoldResult > tileSizes, AffineMap map, ArrayRef< OpFoldResult > lbs, ArrayRef< OpFoldResult > ubs, ArrayRef< OpFoldResult > subShapeSizes, bool omitPartialTileCheck)
Creates an extract_slice/subview op for a single valueToTile with builder.
Definition: Utils.cpp:590
SliceParameters computeSliceParameters(OpBuilder &builder, Location loc, Value valueToTile, ArrayRef< OpFoldResult > tileSizes, AffineMap map, ArrayRef< OpFoldResult > lbs, ArrayRef< OpFoldResult > ubs, ArrayRef< OpFoldResult > subShapeSizes, bool omitPartialTileCheck)
Computes SliceParameters for a single valueToTile assuming that its user is being tiled with the give...
Definition: Utils.cpp:602
SmallVector< Value > ValueVector
An owning vector of values, handy to return from functions.
Definition: SCF.h:70
Include the generated interface declarations.
A struct containing the Linalg producer before and after fusion.
Definition: Utils.h:217
LinalgOp originalProducer
Definition: Utils.h:218
LinalgOp fusedProducer
Definition: Utils.h:219
Utility class used to generate nested loops with ranges described by loopRanges and loop type describ...
Definition: Utils.h:357
static void doit(OpBuilder &b, Location loc, ArrayRef< Range > loopRanges, LinalgOp linalgOp, ArrayRef< utils::IteratorType > iteratorTypes, function_ref< scf::ValueVector(OpBuilder &, Location, ValueRange, ValueRange)> bodyBuilderFn, ArrayRef< linalg::ProcInfo > procInfo={})
Options that allow distribution of loops generated in Linalg transforms to processors while generatin...
Definition: Utils.h:303
ProcInfoCallBackFn procInfo
Callback function that returns the Values for processor ID (procId), and number of processors (nprocs...
Definition: Utils.h:310
Callback function type used to get processor ID, and number of processors used for distribution for a...
Definition: Utils.h:293
DistributionMethod distributionMethod
Definition: Utils.h:296
A struct containing common matchers over linalg op's region.
Definition: Utils.h:328
static std::optional< BinaryOpKind > matchAsScalarBinaryOp(GenericOp op)
Matches the given linalg op if its body is performing binary operation on int or float scalar values ...
Definition: Utils.cpp:96
A struct containg offsets-sizes-strides arguments of the tiled shape.
Definition: Utils.h:140
SmallVector< OpFoldResult > strides
Definition: Utils.h:143
SmallVector< OpFoldResult > sizes
Definition: Utils.h:142
SmallVector< OpFoldResult > offsets
Definition: Utils.h:141