MLIR  22.0.0git
ExtractSliceFromReshapeUtils.cpp
Go to the documentation of this file.
1 //===- ExtractSliceFromReshapeUtils.cpp - Slice reshape rewrites ----------===//
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 file implements rewrites that replace slices of reshape results with
10 // aggregated slices of the reshape source.
11 //
12 //===----------------------------------------------------------------------===//
19 #include "mlir/IR/BuiltinTypes.h"
20 #include "mlir/IR/OpDefinition.h"
21 #include "llvm/ADT/STLExtras.h"
22 
23 using namespace mlir;
24 using namespace mlir::affine;
25 using namespace mlir::tensor;
26 
27 /// A tuple that represents (dimension number, dimension value).
28 using DimAndIndex = std::tuple<unsigned, Value>;
29 
30 /// Transform `dimAndIndex` from the output index space of a (non-rank-reducing)
31 /// slice described by `sliceParams` into the input index space.
33  ArrayRef<Range> sliceParams,
34  const DimAndIndex &dimAndIndex) {
35  AffineExpr d0, s0, s1;
36  bindDims(b.getContext(), d0);
37  bindSymbols(b.getContext(), s0, s1);
38  auto [dim, indexValue] = dimAndIndex;
39  assert(dim < sliceParams.size() && "slice should be non rank-reducing");
40  return std::make_pair(
42  b, loc, s0 + d0 * s1,
43  {indexValue, sliceParams[dim].offset, sliceParams[dim].stride}));
44 }
45 
46 /// Transform `dimAndIndex` from the result tensor index space of a
47 /// CollapseShapeOp to the source tensor index space.
49  OpBuilder &b, Location loc, ArrayRef<ReassociationIndices> reassociation,
50  ArrayRef<OpFoldResult> reshapeSourceShape, const DimAndIndex &dimAndIndex) {
51  const auto &[dim, indexValue] = dimAndIndex;
53  for (int64_t i : reassociation[dim])
54  basis.push_back(reshapeSourceShape[i]);
55  auto delinearized =
56  AffineDelinearizeIndexOp::create(b, loc, indexValue, basis);
57  return delinearized->getResults();
58 }
59 
60 FailureOr<ExtractSliceFromCollapseHelper>
62  OpBuilder &b, tensor::CollapseShapeOp collapseOp,
63  tensor::ExtractSliceOp extractOp) {
64  if (extractOp.getSource().getDefiningOp<tensor::CollapseShapeOp>() !=
65  collapseOp)
66  return failure();
67  SmallVector<Range> ranges;
68  ranges.reserve(extractOp.getSourceType().getRank());
69  for (const auto &[o, s, st] :
70  llvm::zip(extractOp.getMixedOffsets(), extractOp.getMixedSizes(),
71  extractOp.getMixedStrides())) {
72  ranges.push_back({o, s, st});
73  }
74  return ExtractSliceFromCollapseHelper::create(b, collapseOp, ranges);
75 }
76 
77 FailureOr<ExtractSliceFromCollapseHelper>
79  tensor::CollapseShapeOp op,
80  ArrayRef<Range> sliceParams) {
81  // Don't perform this pattern if the collapse op can be simplified by
82  // a rank-reducing extract slice.
83  if (succeeded(mlir::getSimplifyCollapseShapeWithRankReducingSliceInfo(
84  op.getSrcType(), op.getReassociationIndices())))
85  return failure();
86 
87  // Materialize the output shape of the collapse_shape operation. This will
88  // create IR describing the output shape in terms of the input shape.
89  ReifiedRankedShapedTypeDims reifiedShapes;
90  if (failed(reifyResultShapes(b, op, reifiedShapes)))
91  return failure();
92  SmallVector<OpFoldResult> &collapseShapeOutputShape = reifiedShapes[0];
93  SmallVector<ReassociationIndices> reassociationIndices =
94  op.getReassociationIndices();
95 
96  // Determine which of the CollapseShapeOp's result dimensions are sliced
97  // and/or linearized.
98  llvm::SmallBitVector linearizedDimensions =
99  getLinearizedDimensions(reassociationIndices);
100  llvm::SmallBitVector slicedDimensions =
101  getSlicedDimensions(collapseShapeOutputShape, sliceParams);
102 
103  auto collapseShapeInputShape =
104  tensor::getMixedSizes(b, op.getLoc(), op.getSrc());
105 
106  SmallVector<Value> tileSizes;
107  for (unsigned i = 0; i < sliceParams.size(); i++) {
108  if (slicedDimensions[i] && linearizedDimensions[i])
109  tileSizes.push_back(
110  getValueOrCreateConstantIndexOp(b, op.getLoc(), sliceParams[i].size));
111  }
112 
114  op, collapseShapeInputShape, collapseShapeOutputShape, sliceParams,
115  linearizedDimensions, slicedDimensions, tileSizes);
116 }
117 
118 std::pair<Value, SmallVector<Range>>
120  OpBuilder &builder, Location loc, ValueRange tileInductionVars) {
121  // Create the helper class for forming the slice parameters.
122  const SmallVector<ReassociationIndices> reassociationIndices =
123  collapseShapeOp.getReassociationIndices();
124  SliceFromCollapseHelper helper(reassociationIndices, collapseShapeInputShape,
125  collapseShapeOutputShape, sliceParams);
126 
127  // Get the indices of the tiled dims (linearized by the collapse_shape
128  // and sliced by the extract_slice) invert the index spaces
129  // transformations.
130  SmallVector<ValueRange> multiIndices;
131  unsigned loopIdx = 0;
132  for (unsigned i = 0, e = linearizedDimensions.size(); i < e; i++) {
133  if (linearizedDimensions[i] && slicedDimensions[i]) {
134  DimAndIndex tb =
135  invertSliceIndexing(builder, loc, sliceParams,
136  std::make_tuple(i, tileInductionVars[loopIdx++]));
137  multiIndices.push_back(invertCollapseShapeIndexing(
138  builder, loc, reassociationIndices, collapseShapeInputShape, tb));
139  }
140  }
141 
142  SmallVector<Range> extractParams =
143  helper.getExtractSliceParams(builder.getContext(), multiIndices);
144 
145  Value subTileResult = tensor::ExtractSliceOp::create(
146  builder, loc, collapseShapeOp.getSrc(), extractParams);
147 
148  SmallVector<Range> insertParams =
149  helper.getInsertSliceParams(builder.getContext(), tileInductionVars);
150 
151  // Collapse the dimensions of the source slice back down.
152  Value collapsedResult = tensor::CollapseShapeOp::create(
153  builder, loc, subTileResult, reassociationIndices);
154  return std::make_pair(collapsedResult, insertParams);
155 }
156 
157 FailureOr<Operation *>
159  tensor::CollapseShapeOp op, RewriterBase &rewriter) {
160  SmallVector<ReassociationIndices> reassociationIndices =
161  op.getReassociationIndices();
162  RankedTensorType sourceType = op.getSrcType();
163  FailureOr<CollapseShapeRankReducingSliceSimplificationInfo> info =
164  getSimplifyCollapseShapeWithRankReducingSliceInfo(sourceType,
165  reassociationIndices);
166  if (failed(info))
167  return failure();
168 
169  // Create the rank-reducing extract slice op.
170  auto zero = rewriter.getIndexAttr(0);
171  auto one = rewriter.getIndexAttr(1);
172  SmallVector<OpFoldResult> offsets(sourceType.getRank(), zero);
174  tensor::getMixedSizes(rewriter, op.getLoc(), op.getSrc());
175  SmallVector<OpFoldResult> strides(sourceType.getRank(), one);
176  auto sliceOp = tensor::ExtractSliceOp::create(
177  rewriter, op.getLoc(), info->sliceResultType, op.getSrc(), offsets, sizes,
178  strides);
179 
180  if (!info->newReassociationIndices.has_value()) {
181  rewriter.replaceOp(op, sliceOp.getResult());
182  return sliceOp.getOperation();
183  }
184 
185  return rewriter
186  .replaceOpWithNewOp<tensor::CollapseShapeOp>(
187  op, sliceOp.getResult(), *info->newReassociationIndices)
188  .getOperation();
189 }
static ValueRange invertCollapseShapeIndexing(OpBuilder &b, Location loc, ArrayRef< ReassociationIndices > reassociation, ArrayRef< OpFoldResult > reshapeSourceShape, const DimAndIndex &dimAndIndex)
Transform dimAndIndex from the result tensor index space of a CollapseShapeOp to the source tensor in...
static DimAndIndex invertSliceIndexing(OpBuilder &b, Location loc, ArrayRef< Range > sliceParams, const DimAndIndex &dimAndIndex)
Transform dimAndIndex from the output index space of a (non-rank-reducing) slice described by slicePa...
std::tuple< unsigned, Value > DimAndIndex
A tuple that represents (dimension number, dimension value).
Base type for affine expression.
Definition: AffineExpr.h:68
IntegerAttr getIndexAttr(int64_t value)
Definition: Builders.cpp:103
MLIRContext * getContext() const
Definition: Builders.h:55
This class defines the main interface for locations in MLIR and acts as a non-nullable wrapper around...
Definition: Location.h:76
This class helps build Operations.
Definition: Builders.h:205
This class coordinates the application of a rewrite on a set of IR, providing a way for clients to tr...
Definition: PatternMatch.h:358
virtual void replaceOp(Operation *op, ValueRange newValues)
Replace the results of the given (original) operation with the specified list of values (replacements...
OpTy replaceOpWithNewOp(Operation *op, Args &&...args)
Replace the results of the given (original) op with a new op that is created without verification (re...
Definition: PatternMatch.h:519
This class provides an abstraction over the different types of ranges over Values.
Definition: ValueRange.h:387
This class represents an instance of an SSA value in the MLIR system, representing a computable value...
Definition: Value.h:96
This class assists with generating IR required to materialize an arbitrary-sized slice from the resul...
std::pair< Value, SmallVector< Range > > emitLoopNestBody(OpBuilder &builder, Location loc, ValueRange tileInductionVars)
Generates the IR inside of the caller's loop nest for 1) inverting the index mappings of the ExtractS...
static FailureOr< ExtractSliceFromCollapseHelper > create(OpBuilder &b, tensor::CollapseShapeOp op, ArrayRef< Range > sliceParams)
Given a CollapseShapeOp and a set of ranges describing the desired slice of its result,...
AffineApplyOp makeComposedAffineApply(OpBuilder &b, Location loc, AffineMap map, ArrayRef< OpFoldResult > operands, bool composeAffineMin=false)
Returns a composed AffineApplyOp by composing map and operands with other AffineApplyOps supplying th...
Definition: AffineOps.cpp:1278
FailureOr< Operation * > simplifyCollapseShapeWithRankReducingExtractSlice(tensor::CollapseShapeOp op, RewriterBase &rewriter)
Tries to simplify a tensor.collapse_shape operation by inserting a single rank-reducing tensor....
SmallVector< OpFoldResult > getMixedSizes(OpBuilder &builder, Location loc, Value value)
Return the dimensions of the given tensor value.
Definition: TensorOps.cpp:70
Include the generated interface declarations.
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...
LogicalResult reifyResultShapes(OpBuilder &b, Operation *op, ReifiedRankedShapedTypeDims &reifiedReturnShapes)
Reify the shape of the result of an operation (typically in terms of the shape of its operands).
void bindDims(MLIRContext *ctx, AffineExprTy &...exprs)
Bind a list of AffineExpr references to DimExpr at positions: [0 .
Definition: AffineExpr.h:311
void bindSymbols(MLIRContext *ctx, AffineExprTy &...exprs)
Bind a list of AffineExpr references to SymbolExpr at positions: [0 .
Definition: AffineExpr.h:325
Value getValueOrCreateConstantIndexOp(OpBuilder &b, Location loc, OpFoldResult ofr)
Converts an OpFoldResult to a Value.
Definition: Utils.cpp:112
llvm::SmallBitVector getLinearizedDimensions(ArrayRef< ReassociationIndices > reassociationIndices)
Determine which dimensions are linearized by a tensor.collapse_shape op by inspecting its reassociati...