MLIR  18.0.0git
ReshapeOpsUtils.h
Go to the documentation of this file.
1 //===- ReshapeOpsUtils.h - Utilities used by reshape ops --*- 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 // This header file defines utilities and common canonicalization patterns for
10 // reshape operations.
11 //
12 //===----------------------------------------------------------------------===//
13 
14 #ifndef MLIR_DIALECT_UTILS_RESHAPEOPSUTILS_H
15 #define MLIR_DIALECT_UTILS_RESHAPEOPSUTILS_H
16 
19 #include "mlir/IR/PatternMatch.h"
20 #include "mlir/Support/LLVM.h"
21 #include "llvm/ADT/StringRef.h"
22 #include <optional>
23 
24 namespace mlir {
25 
29 
30 /// Attribute name for the ArrayAttr which encodes reassociation indices.
31 constexpr StringRef getReassociationAttrName() { return "reassociation"; }
32 
33 /// Compose reassociation maps that are used in pair of reshape ops where one
34 /// is a producer and other is the consumer. Only valid to use this method when
35 /// both the producer and consumer are collapsing dimensions or both are
36 /// expanding dimensions.
37 ///
38 /// For example,
39 /// producerReassociation = [[0, 1], [2], [3, 4]]
40 /// consumerReassociation = [[0, 1], [2]]
41 ///
42 /// is folded into
43 ///
44 /// result = [[0, 1, 2], [3, 4]].
45 std::optional<SmallVector<ReassociationIndices>> composeReassociationIndices(
46  ArrayRef<ReassociationIndices> producerReassociations,
47  ArrayRef<ReassociationIndices> consumerReassociations,
48  MLIRContext *context);
49 
50 /// Convert reassociation indices to affine expressions.
51 SmallVector<SmallVector<AffineExpr, 2>, 2> convertReassociationIndicesToExprs(
52  MLIRContext *context, ArrayRef<ReassociationIndices> reassociationIndices);
53 
54 /// Constructs affine maps out of Array<Array<AffineExpr>>.
55 SmallVector<AffineMap, 4>
56 getSymbolLessAffineMaps(ArrayRef<ReassociationExprs> reassociation);
57 
58 /// Wraps a list of reassociations in an ArrayAttr.
59 ArrayAttr
61  ArrayRef<ReassociationIndices> reassociation);
62 
63 /// Convert Array<Array<AffineExpr>> to Array<Array<int64_t>>.
64 SmallVector<ReassociationIndices, 2> convertReassociationMapsToIndices(
65  OpBuilder &b, ArrayRef<ReassociationExprs> reassociationExprs);
66 
67 /// Return the reassociations maps to use to reshape given the source type and
68 /// the target type when possible. Return std::nullopt when this computation
69 /// failed.
70 std::optional<SmallVector<ReassociationIndices>>
71 getReassociationIndicesForReshape(ShapedType sourceType, ShapedType targetType);
72 
73 /// Returns the reassociation maps to collapse `sourceShape` to `targetShape` if
74 /// possible.
75 std::optional<SmallVector<ReassociationIndices>>
76 getReassociationIndicesForCollapse(ArrayRef<int64_t> sourceShape,
77  ArrayRef<int64_t> targetShape);
78 
79 /// Return true if the reassociation specification is valid, false otherwise.
80 /// When false, the `invalidIndex` integer pointer is optionally filled with the
81 /// index of the offending reassociation map.
82 bool isReassociationValid(ArrayRef<AffineMap> reassociation,
83  int *invalidIndex = nullptr);
84 
85 template <typename ReshapeOpTy, typename InverseReshapeOpTy>
86 static OpFoldResult foldReshapeOp(ReshapeOpTy reshapeOp,
87  ArrayRef<Attribute> operands) {
88  // Fold producer-consumer reshape ops that where the operand type of the
89  // producer is same as the return type of the consumer.
90  auto reshapeSrcOp =
91  reshapeOp.getSrc().template getDefiningOp<InverseReshapeOpTy>();
92  if (reshapeSrcOp && reshapeSrcOp.getSrcType() == reshapeOp.getResultType())
93  return reshapeSrcOp.getSrc();
94  // Reshape of a constant can be replaced with a new constant.
95  if (auto elements = dyn_cast_or_null<DenseElementsAttr>(operands.front())) {
96  return elements.reshape(cast<ShapedType>(reshapeOp.getResult().getType()));
97  }
98  return nullptr;
99 }
100 
101 /// Common verifier for reshape-like types. Fills `expandedType` and
102 ///`collapsedType` with the proper `src` or `result` type.
103 template <typename Op, typename T>
104 static LogicalResult verifyReshapeLikeTypes(Op op, T expandedType,
105  T collapsedType, bool isExpansion) {
106  unsigned expandedRank = expandedType.getRank();
107  unsigned collapsedRank = collapsedType.getRank();
108  if (expandedRank < collapsedRank)
109  return op.emitOpError("expected the type ")
110  << expandedType
111  << " to have higher rank than the type = " << collapsedType;
112  if (expandedRank == 0)
113  return op.emitOpError("expected non-zero memref ranks");
114  if (expandedRank == collapsedRank)
115  return op.emitOpError("expected to collapse or expand dims");
116 
117  if (collapsedRank == 0) {
118  // If collapsed rank is 0, then expanded type must be static shaped and of
119  // sizes 1.
120  if (llvm::any_of(expandedType.getShape(),
121  [](int64_t dim) -> bool { return dim != 1; }))
122  return op.emitOpError("invalid to reshape tensor/memref with non-unit "
123  "extent dimensions to zero-rank tensor/memref");
124  return success();
125  }
126  if (collapsedRank != op.getReassociation().size())
127  return op.emitOpError("expected rank of the collapsed type(")
128  << collapsedRank << ") to be the number of reassociation maps("
129  << op.getReassociation().size() << ")";
130  auto maps = op.getReassociationMaps();
131  for (auto it : llvm::enumerate(maps))
132  if (it.value().getNumDims() != expandedRank)
133  return op.emitOpError("expected reassociation map #")
134  << it.index() << " of same rank as expanded memref("
135  << expandedRank << "), but got " << it.value().getNumDims();
136  int invalidIdx = 0;
137  if (!isReassociationValid(maps, &invalidIdx))
138  return op.emitOpError("expected reassociation map #")
139  << invalidIdx << " to be valid and contiguous";
140  return verifyReshapeLikeShapes(op, collapsedType, expandedType, isExpansion);
141 }
142 
143 /// Verify that shapes of the reshaped types using following rules
144 /// 1) if a dimension in the collapsed type is static, then the corresponding
145 /// dimensions in the expanded shape should be
146 /// a) static
147 /// b) the product should be same as the collaped shape.
148 /// 2) if a dimension in the collaped type is dynamic, one and only one of the
149 /// corresponding dimensions in the expanded type should be dynamic. This
150 /// rule is only needed with reshape operations that are expanding.
151 LogicalResult reshapeLikeShapesAreCompatible(
152  function_ref<LogicalResult(const Twine &)> emitError,
153  ArrayRef<int64_t> collapsedShape, ArrayRef<int64_t> expandedShape,
154  ArrayRef<ReassociationIndices> reassociationMaps, bool isExpandingReshape);
155 
156 template <typename OpTy>
157 static LogicalResult verifyReshapeLikeShapes(OpTy op, ShapedType collapsedType,
158  ShapedType expandedType,
159  bool isExpandingReshape) {
161  [&](const Twine &msg) { return op->emitOpError(msg); },
162  collapsedType.getShape(), expandedType.getShape(),
163  op.getReassociationIndices(), isExpandingReshape);
164 }
165 
166 /// Returns true iff the type is a MemRefType and has a non-identity layout.
167 bool hasNonIdentityLayout(Type type);
168 
169 /// Pattern to collapse producer/consumer reshape ops that are both collapsing
170 /// dimensions or are both expanding dimensions.
171 template <typename ReshapeOpTy>
172 struct ComposeReassociativeReshapeOps : public OpRewritePattern<ReshapeOpTy> {
174  LogicalResult matchAndRewrite(ReshapeOpTy reshapeOp,
175  PatternRewriter &rewriter) const override {
176  auto srcReshapeOp =
177  reshapeOp.getSrc().template getDefiningOp<ReshapeOpTy>();
178  if (!srcReshapeOp)
179  return failure();
180 
181  ShapedType resultType = reshapeOp.getResultType();
182 
183  if (hasNonIdentityLayout(srcReshapeOp.getSrc().getType()) ||
184  hasNonIdentityLayout(reshapeOp.getSrc().getType()) ||
185  hasNonIdentityLayout(reshapeOp.getResult().getType()))
186  return failure();
187 
188  std::optional<SmallVector<ReassociationIndices>> reassociationIndices =
189  composeReassociationIndices(srcReshapeOp.getReassociationIndices(),
190  reshapeOp.getReassociationIndices(),
191  rewriter.getContext());
192  if (!reassociationIndices)
193  return failure();
194  rewriter.replaceOpWithNewOp<ReshapeOpTy>(
195  reshapeOp, resultType, srcReshapeOp.getSrc(), *reassociationIndices);
196  return success();
197  }
198 };
199 
200 /// Pattern to compose
201 /// `collapse_shape(expand_shape(%src, reassociation_1), reassociation_2)`.
202 /// In that case both `srcType` and `resultType` can be expressed as a function
203 /// of `intermediateType`.
204 /// In order to demonstrate the approach, let's assume that `rank(srcType) >
205 /// `rank(resultType)`, i.e. the resulting operation should be `collapse_shape`.
206 /// In that case, we can iterate over every set of indices in `reassociation_2`
207 /// and try to find ids of sets of indices in `reassociation_1` that cover it
208 /// completely.
209 ///
210 /// Example:
211 ///
212 /// %0 = tensor.expand_shape %arg [[0], [1], [2, 3]]
213 /// : tensor<?x?x?xi64> into tensor<?x?x?x1xi64>
214 /// %1 = tensor.collapse_shape %0 [[0, 1], [2, 3]]
215 /// : tensor<?x?x?x1xi64> into tensor<?x?xi64>
216 ///
217 /// can be canonicalized into
218 ///
219 /// %0 = tensor.collapse_shape %arg [[0, 1], [2]]
220 /// : tensor<?x?x?xi64> into tensor<?x?xi64>
221 ///
222 /// because [0] and [1] from `expand_shape` reassociation cover completely
223 /// `[0, 1]` from `collapse_shape`. If it is impossible to find such union of
224 /// indices, then we fail.
225 //
226 /// When `rank(srcType) < rank(resultType)`, then we just swap `reassociation_1`
227 /// `reassociation_2` and produce `expand_shape`.
228 template <typename CollapseOpTy, typename ExpandOpTy, typename CastOpTy>
229 struct ComposeCollapseOfExpandOp : public OpRewritePattern<CollapseOpTy> {
231  LogicalResult matchAndRewrite(CollapseOpTy collapseOp,
232  PatternRewriter &rewriter) const override {
233  auto expandOp = collapseOp.getSrc().template getDefiningOp<ExpandOpTy>();
234  if (!expandOp)
235  return failure();
236 
237  ShapedType srcType = expandOp.getSrcType();
238  ShapedType resultType = collapseOp.getResultType();
239 
240  if (hasNonIdentityLayout(collapseOp.getSrc().getType()) ||
241  hasNonIdentityLayout(expandOp.getSrc().getType()) ||
242  hasNonIdentityLayout(expandOp.getResult().getType()))
243  return failure();
244 
245  int64_t srcRank = srcType.getRank();
246  int64_t resultRank = resultType.getRank();
247  if (srcType == resultType)
248  return failure();
249 
250  SmallVector<ReassociationIndices, 4> higherRankReassociation,
251  lowerRankReassociation;
252 
253  if (srcRank > resultRank) {
254  higherRankReassociation = expandOp.getReassociationIndices();
255  lowerRankReassociation = collapseOp.getReassociationIndices();
256  } else {
257  higherRankReassociation = collapseOp.getReassociationIndices();
258  lowerRankReassociation = expandOp.getReassociationIndices();
259  }
260 
261  size_t higherRankIndicesID = 0;
262  SmallVector<ReassociationIndices, 4> composedReassociation;
263  for (const auto &lowerRankIndices : lowerRankReassociation) {
264  ReassociationIndices composedIndices;
265  while (higherRankIndicesID < higherRankReassociation.size()) {
266  auto rightmostIndex =
267  higherRankReassociation[higherRankIndicesID].back();
268  if (rightmostIndex > lowerRankIndices.back())
269  return failure();
270  composedIndices.push_back(higherRankIndicesID++);
271  if (rightmostIndex == lowerRankIndices.back())
272  break;
273  }
274  composedReassociation.push_back(composedIndices);
275  }
276  if (srcRank > resultRank) {
277  rewriter.replaceOpWithNewOp<CollapseOpTy>(
278  collapseOp, resultType, expandOp.getSrc(), composedReassociation);
279  } else if (srcRank < resultRank) {
280  rewriter.replaceOpWithNewOp<ExpandOpTy>(
281  collapseOp, resultType, expandOp.getSrc(), composedReassociation);
282  } else {
283  // Collapses/expansions that do not change the rank are not allowed. Use
284  // a cast instead.
285  assert(llvm::equal(srcType.getShape(), resultType.getShape()) &&
286  "expected same shape");
287  rewriter.replaceOpWithNewOp<CastOpTy>(collapseOp, resultType,
288  expandOp.getSrc());
289  }
290  return success();
291  }
292 };
293 
294 template <typename ExpandOpTy, typename CollapseOpTy>
295 struct ComposeExpandOfCollapseOp : public OpRewritePattern<ExpandOpTy> {
297  LogicalResult matchAndRewrite(ExpandOpTy expandOp,
298  PatternRewriter &rewriter) const override {
299  auto collapseOp = expandOp.getSrc().template getDefiningOp<CollapseOpTy>();
300  if (!collapseOp)
301  return failure();
302 
303  ShapedType srcType = collapseOp.getSrcType();
304  ShapedType resultType = expandOp.getResultType();
305 
306  if (hasNonIdentityLayout(expandOp.getSrc().getType()) ||
307  hasNonIdentityLayout(collapseOp.getSrc().getType()) ||
308  hasNonIdentityLayout(collapseOp.getResult().getType()))
309  return failure();
310 
311  int64_t srcRank = srcType.getRank();
312  int64_t resultRank = resultType.getRank();
313  if (srcType == resultType)
314  return failure();
315 
316  auto srcReassociation = collapseOp.getReassociationIndices();
317  auto resultReassociation = expandOp.getReassociationIndices();
318  if (srcRank > resultRank) {
319  auto composedReassociation = findCollapsingReassociation(
320  srcReassociation, resultReassociation, srcType.getShape(),
321  resultType.getShape());
322  if (!composedReassociation)
323  return failure();
324 
325  rewriter.replaceOpWithNewOp<CollapseOpTy>(
326  expandOp, resultType, collapseOp.getSrc(), *composedReassociation);
327  return success();
328  }
329  auto composedReassociation =
330  findCollapsingReassociation(resultReassociation, srcReassociation,
331  resultType.getShape(), srcType.getShape());
332  if (!composedReassociation)
333  return failure();
334 
335  rewriter.replaceOpWithNewOp<ExpandOpTy>(
336  expandOp, resultType, collapseOp.getSrc(), *composedReassociation);
337  return success();
338  }
339 
340 private:
341  // Attempts to find a way to collapse `srcShape` to `resultShape` by
342  // collapsing subshapes defined by the reassociation indices.
343  std::optional<SmallVector<ReassociationIndices>> findCollapsingReassociation(
344  ArrayRef<ReassociationIndices> srcReassociation,
345  ArrayRef<ReassociationIndices> resultReassociation,
346  ArrayRef<int64_t> srcShape, ArrayRef<int64_t> resultShape) const {
347  SmallVector<ReassociationIndices, 4> composedReassociation;
348 
349  if (srcReassociation.empty())
350  return {getReassociationIndicesForCollapse(srcShape, resultShape)};
351 
352  for (auto item : llvm::zip(srcReassociation, resultReassociation)) {
353  auto &srcIndices = std::get<0>(item);
354  auto &resultIndices = std::get<1>(item);
355  auto srcSubShape = srcShape.slice(srcIndices.front(), srcIndices.size());
356  auto resultSubShape =
357  resultShape.slice(resultIndices.front(), resultIndices.size());
358 
359  if (srcSubShape.size() == resultSubShape.size()) {
360  if (srcSubShape == resultSubShape)
361  composedReassociation.push_back(srcIndices);
362  else
363  return std::nullopt;
364  }
365 
366  // Find reassociation to collapse `srcSubShape` into `resultSubShape`.
367  auto subShapeReassociation =
368  getReassociationIndicesForCollapse(srcSubShape, resultSubShape);
369  if (!subShapeReassociation)
370  return std::nullopt;
371 
372  // Remap the subshape indices back to the original srcShape.
373  for (auto &subshape_indices : *subShapeReassociation) {
374  ReassociationIndices shape_indices;
375  for (int64_t index : subshape_indices)
376  shape_indices.push_back(srcIndices.front() + index);
377  composedReassociation.push_back(shape_indices);
378  }
379  }
380  return {std::move(composedReassociation)};
381  }
382 };
383 
384 /// The input parameters `offsets`, `sizes`, `strides` specify a rectangular
385 /// non rank-reducing slice of the collapse_shape output. Try to find which
386 /// dimensions have been sliced and which dimensions are not sliced (offset = 0,
387 /// size = dim, size = 1). Note that this conservative as it cannot detect if a
388 /// dynamic size corresponds to the full tensor dimension or not.
389 llvm::SmallBitVector getSlicedDimensions(ArrayRef<OpFoldResult> sliceInputShape,
390  ArrayRef<Range> sliceParams);
391 
392 /// Determine which dimensions are linearized by a `tensor.collapse_shape` op by
393 /// inspecting its reassociation indices.
394 llvm::SmallBitVector
395 getLinearizedDimensions(ArrayRef<ReassociationIndices> reassociationIndices);
396 
397 /// Given the parameters for both operations in a `CollapseShape->ExtractSlice`
398 /// chain and reified source and result shapes of the CollapseShapeOp, this
399 /// class provides two functions that assist with directly forming the result
400 /// of the extract slice by "tiling the CollapseShapeOp by 1".
401 //// Example:
402 // clang-format off
403 /// ```
404 /// %0 = linalg.generic ... -> tensor<3x7x11x10xf32>
405 /// %1 = tensor.collapse_shape %0 [[0, 1, 2], [3]] : ... to tensor<341x10xf32>
406 /// %2 = tensor.extract_slice %1 [13, 0] [10, 10] [2, 1] : .... tensor<10x10xf32>
407 /// ```
408 /// This class helps build the below IR to replace %2:
409 /// ```
410 /// %dest = tensor.empty() : tensor<10x10xf32>
411 /// %2 = scf.for %iv = %c0 to %c10 step %c1 iter_args(%arg0) -> tensor<10x10xf32> {
412 /// %linear_index = affine.apply affine_map<(d0)[]->(d0*2 + 11)>(%iv)
413 /// %3:3 = arith.delinearize_index %iv into (3, 7, 11)
414 ///
415 /// // This function takes %3 (multiIndices) and the parameters for the slice below.
416 /// %4 = tensor.extract_slice %0 [%3#0, %3#1, %3#2, 0] [1, 1, 1, 10] [1, 1, 1, 1] :
417 /// tensor<3x7x11x10xf32> to tensor<1x1x1x10xf32>
418 ///
419 /// %5 = tensor.collapse_shape %4 [[0, 1, 2], [3]] :
420 /// tensor<1x1x1x10xf32> into tensor<1x10xf32>
421 /// %6 = tensor.insert_slice %5 into %arg0 [%iv, 0] [1, 10] [1, 1] :
422 /// tensor<1x10xf32> into tensor<10x10xf32>
423 /// scf.yield %6 : tensor<10x10xf32>
424 /// }
425 /// ```
426 // clang-format on
427 class SliceFromCollapseHelper {
428 public:
429  SliceFromCollapseHelper(ArrayRef<ReassociationIndices> reassociationIndices,
430  ArrayRef<OpFoldResult> collapseShapeInputShape,
431  ArrayRef<OpFoldResult> collapseShapeOutputShape,
432  ArrayRef<Range> extractSliceParams)
433  : reassociationIndices(reassociationIndices),
434  collapseShapeInputShape(collapseShapeInputShape),
435  collapseShapeOutputShape(collapseShapeOutputShape),
436  sliceParams(extractSliceParams),
437  linearizedDimensions(getLinearizedDimensions(reassociationIndices)),
438  slicedDimensions(getSlicedDimensions(collapseShapeOutputShape,
439  extractSliceParams)) {}
440 
441  /// This function takes multi-indices and maps them to ExtractSlice parameters
442  /// in the index space of the CollapseShape's source tensor. This function's
443  /// signature can be described by `(D_0, D_1,.. D_{n-1}) -> (offsets, sizes,
444  /// strides)` where `n` the number of "tiled dimensions", which are the
445  /// dimensions of the output that are linearized by the collapse shape op and
446  /// are also sliced. Each `D_i` is a tuple that must represent a valid
447  /// multi-index for the `i-th` tiled dimension. In the example above, there is
448  /// only one tiled dimension (D_0) and `arith.delinearize_index` produces the
449  /// multi-index (%3) that would be passed to this function to generate the
450  /// parameters for the `tensor.extract_slice` op (%4).
451  SmallVector<Range> getExtractSliceParams(MLIRContext *ctx,
452  ArrayRef<ValueRange> multiIndices);
453 
454  /// This function takes indices in the index space of the "tiled dimensions"
455  /// described above and returns a set of Range variables that describe how the
456  /// slice should be inserted into the destination. In the example above, `%iv`
457  /// would be passed to this function to generate the parameters for the
458  /// `tensor.insert_slice` op producing %6.
459  SmallVector<Range> getInsertSliceParams(MLIRContext *ctx,
460  ValueRange tileIndices);
461 
462 private:
463  SmallVector<ReassociationIndices> reassociationIndices;
464  SmallVector<OpFoldResult> collapseShapeInputShape;
465  SmallVector<OpFoldResult> collapseShapeOutputShape;
466  SmallVector<Range> sliceParams;
467  llvm::SmallBitVector linearizedDimensions;
468  llvm::SmallBitVector slicedDimensions;
469 };
470 
471 /// Parameters required to simplify a collapsing reshape op with a rank-reducing
472 /// slice operation. See `getSimplifyCollapseShapeWithRankReducingSliceInfo`.
473 struct CollapseShapeRankReducingSliceSimplificationInfo {
474  /// The shape of the output of the rank-reducing slice.
475  RankedTensorType sliceResultType;
476  /// The reassociation indices for the new collapse shape op, if required. If
477  /// `std::nullopt`, the slice should replace the collapse shape op.
478  std::optional<SmallVector<ReassociationIndices>> newReassociationIndices;
479 };
480 
481 /// A collapsing reshape operation can sometimes be simplified or eliminated by
482 /// inserting a single rank-reducing slice operation between it and the source
483 /// tensor. The slice op will either take the place of the source, allowing for
484 /// a new, simpler reshape op to replace the original, or the reshape op will be
485 /// completely replaced by the slice result.
486 ///
487 /// This function returns the parameters required to implement this pattern. If
488 /// the pattern is not applicable, then failure is returned.
489 ///
490 /// ### Example:
491 /// ```
492 /// %result = tensor.collapse_shape %0 [[0, 1], [2, 3]]
493 /// : tensor<?x1x30x10xf32> to tensor<?x300xf32>
494 /// ```
495 /// can be transformed to
496 /// ```
497 /// %tmp = tensor.extract_slice %0 [0, 0, 0, 0]
498 /// [0, %dim1, 30, 30]
499 /// [1, 1, 1 1]
500 /// : tensor<?x1x30x10xf32> to tensor<?x30x10xf32>
501 /// %result = tensor.collapse_shape %tmp [[0], [1, 2]]
502 /// : tensor<?x30x10xf32> to tensor<?x300xf32>
503 /// ```
504 ///
505 /// ### Example:
506 /// ```
507 /// %result = tensor.collapse_shape %1 [[0, 1], [2]]
508 /// : tensor<?x1x30xf32> to tensor<?x30xf32>
509 /// ```
510 /// can be transformed to
511 /// ```
512 /// %result = tensor.extract_slice %1 [0, 0, 0]
513 /// [%dim2, 1, 30]
514 /// [1, 1, 1]
515 /// : tensor<?x1x30xf32> to tensor<?x30xf32>
516 /// ```
517 FailureOr<CollapseShapeRankReducingSliceSimplificationInfo>
518 getSimplifyCollapseShapeWithRankReducingSliceInfo(
519  RankedTensorType sourceType,
520  ArrayRef<ReassociationIndices> reassociationIndices);
521 
522 struct PackingMetadata {
523  SmallVector<int64_t> insertPositions;
524  SmallVector<int64_t> outerPositions;
525  SmallVector<ReassociationIndices> reassociations;
526 };
527 
528 /// Given a vector of `positions` indices representing desired packing insertion
529 /// points into a target vector (i.e. pack/unpack.inner_dim_pos), compute the
530 /// final positions in the target shape as well as the reshape reassociations.
531 // Note: This should not be called with a large positions array (or the
532 // implementation needs to be updated to use an N.log N sort instead of
533 // repeated N^2 counts).
534 PackingMetadata computePackingMetadata(int64_t packedRank,
535  ArrayRef<int64_t> innerDimPos);
536 } // namespace mlir
537 
538 #endif // MLIR_DIALECT_UTILS_RESHAPEOPSUTILS_H
MLIRContext * getContext() const
Definition: Builders.h:55
This class represents a single result from folding an operation.
Definition: OpDefinition.h:266
This provides public APIs that all operations should have.
InFlightDiagnostic emitOpError(const Twine &message={})
Emit an error with the op name prefixed, like "'dim' op " which is convenient for verifiers.
Definition: Operation.cpp:640
A special type of RewriterBase that coordinates the application of a rewrite pattern on the current I...
Definition: PatternMatch.h:727
OpTy replaceOpWithNewOp(Operation *op, Args &&...args)
Replaces the result op with a new op that is created without verification.
Definition: PatternMatch.h:539
constexpr void enumerate(std::tuple< Tys... > &tuple, CallbackT &&callback)
Definition: Matchers.h:285
Include the generated interface declarations.
LogicalResult failure(bool isFailure=true)
Utility function to generate a LogicalResult.
Definition: LogicalResult.h:62
llvm::function_ref< Fn > function_ref
Definition: LLVM.h:147
llvm::SmallBitVector getSlicedDimensions(ArrayRef< OpFoldResult > sliceInputShape, ArrayRef< Range > sliceParams)
The input parameters offsets, sizes, strides specify a rectangular non rank-reducing slice of the col...
constexpr StringRef getReassociationAttrName()
Attribute name for the ArrayAttr which encodes reassociation indices.
bool hasNonIdentityLayout(Type type)
Returns true iff the type is a MemRefType and has a non-identity layout.
static OpFoldResult foldReshapeOp(ReshapeOpTy reshapeOp, ArrayRef< Attribute > operands)
InFlightDiagnostic emitError(Location loc)
Utility method to emit an error message using this location.
SmallVector< AffineMap, 4 > getSymbolLessAffineMaps(ArrayRef< ReassociationExprs > reassociation)
Constructs affine maps out of Array<Array<AffineExpr>>.
static LogicalResult verifyReshapeLikeTypes(Op op, T expandedType, T collapsedType, bool isExpansion)
Common verifier for reshape-like types.
LogicalResult success(bool isSuccess=true)
Utility function to generate a LogicalResult.
Definition: LogicalResult.h:56
LogicalResult reshapeLikeShapesAreCompatible(function_ref< LogicalResult(const Twine &)> emitError, ArrayRef< int64_t > collapsedShape, ArrayRef< int64_t > expandedShape, ArrayRef< ReassociationIndices > reassociationMaps, bool isExpandingReshape)
Verify that shapes of the reshaped types using following rules 1) if a dimension in the collapsed typ...
std::optional< SmallVector< ReassociationIndices > > getReassociationIndicesForReshape(ShapedType sourceType, ShapedType targetType)
Return the reassociations maps to use to reshape given the source type and the target type when possi...
std::optional< SmallVector< ReassociationIndices > > getReassociationIndicesForCollapse(ArrayRef< int64_t > sourceShape, ArrayRef< int64_t > targetShape)
Returns the reassociation maps to collapse sourceShape to targetShape if possible.
ArrayAttr getReassociationIndicesAttribute(OpBuilder &b, ArrayRef< ReassociationIndices > reassociation)
Wraps a list of reassociations in an ArrayAttr.
SmallVector< SmallVector< AffineExpr, 2 >, 2 > convertReassociationIndicesToExprs(MLIRContext *context, ArrayRef< ReassociationIndices > reassociationIndices)
Convert reassociation indices to affine expressions.
bool isReassociationValid(ArrayRef< AffineMap > reassociation, int *invalidIndex=nullptr)
Return true if the reassociation specification is valid, false otherwise.
std::optional< SmallVector< ReassociationIndices > > composeReassociationIndices(ArrayRef< ReassociationIndices > producerReassociations, ArrayRef< ReassociationIndices > consumerReassociations, MLIRContext *context)
Compose reassociation maps that are used in pair of reshape ops where one is a producer and other is ...
SmallVector< int64_t, 2 > ReassociationIndices
llvm::SmallBitVector getLinearizedDimensions(ArrayRef< ReassociationIndices > reassociationIndices)
Determine which dimensions are linearized by a tensor.collapse_shape op by inspecting its reassociati...
static LogicalResult verifyReshapeLikeShapes(OpTy op, ShapedType collapsedType, ShapedType expandedType, bool isExpandingReshape)
SmallVector< ReassociationIndices, 2 > convertReassociationMapsToIndices(OpBuilder &b, ArrayRef< ReassociationExprs > reassociationExprs)
Convert Array<Array<AffineExpr>> to Array<Array<int64_t>>.
Pattern to compose collapse_shape(expand_shape(src, reassociation_1), reassociation_2).
LogicalResult matchAndRewrite(CollapseOpTy collapseOp, PatternRewriter &rewriter) const override
LogicalResult matchAndRewrite(ExpandOpTy expandOp, PatternRewriter &rewriter) const override
Pattern to collapse producer/consumer reshape ops that are both collapsing dimensions or are both exp...
LogicalResult matchAndRewrite(ReshapeOpTy reshapeOp, PatternRewriter &rewriter) const override
This class represents an efficient way to signal success or failure.
Definition: LogicalResult.h:26
OpRewritePattern is a wrapper around RewritePattern that allows for matching and rewriting against an...
Definition: PatternMatch.h:357