MLIR  22.0.0git
ReifyResultShapes.cpp
Go to the documentation of this file.
1 //===- ReifyResultShapes.cpp - Reify result shapes ------------------------===//
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 transform reifies result shapes of `ReifyRankedShapedTypeOpInterface`
10 // operations with ranked `memref` and `tensor` results.
11 //
12 //===----------------------------------------------------------------------===//
13 
15 
21 #include "llvm/Support/InterleavedRange.h"
22 
23 #define DEBUG_TYPE "reify-result-shapes"
24 #define DBGS() (llvm::dbgs() << "[" DEBUG_TYPE << "]: ")
25 
26 namespace mlir {
27 namespace memref {
28 #define GEN_PASS_DEF_REIFYRESULTSHAPESPASS
29 #include "mlir/Dialect/MemRef/Transforms/Passes.h.inc"
30 } // namespace memref
31 } // namespace mlir
32 
33 using namespace mlir;
34 
35 /// Reifies the results of `op`, potentially replacing `op` with a reified
36 /// version. Returns `failure` if `mlir::reifyResultShapes` returned failure,
37 /// otherwise it always succeeds. Users of this transform should always expect
38 /// it to modify the IR, even when it fails. If any of the result types changes,
39 /// the transform will insert cast operations to the old type to keep the IR
40 /// consistent.
41 static LogicalResult reifyOpResultShapes(RewriterBase &rewriter,
42  ReifyRankedShapedTypeOpInterface op) {
43  LLVM_DEBUG({ DBGS() << " reifying op: " << op << "\n"; });
44  // Get the reified out shapes.
45  ReifiedRankedShapedTypeDims reifiedResultShapes;
46  if (failed(mlir::reifyResultShapes(rewriter, op, reifiedResultShapes)) ||
47  reifiedResultShapes.empty()) {
48  return op->emitWarning() << "failed to get the reified shapes";
49  }
50 
51  bool modified = false;
52  // Compute the new output types.
53  SmallVector<Type> outTypes;
54  for (const auto &[oldTy, reifiedShape] :
55  llvm::zip(op->getResultTypes(), reifiedResultShapes)) {
56  // Skip if it's not a memref or tensor type.
57  if (!isa<RankedTensorType, MemRefType>(oldTy)) {
58  outTypes.push_back(oldTy);
59  continue;
60  }
61 
62  ShapedType shapedTy = dyn_cast<ShapedType>(oldTy);
63 
64  SmallVector<int64_t> shape = llvm::to_vector(shapedTy.getShape());
65  for (auto &&[dim, ofr] : llvm::zip_equal(shape, reifiedShape)) {
66  std::optional<int64_t> maybeCst = getConstantIntValue(ofr);
67  // If the reified dim is dynamic set it appropriately.
68  if (!maybeCst.has_value()) {
69  dim = ShapedType::kDynamic;
70  continue;
71  }
72  // Set the static dim.
73  dim = *maybeCst;
74  }
75 
76  // If the shape didn't change continue.
77  if (shape == shapedTy.getShape()) {
78  outTypes.push_back(oldTy);
79  continue;
80  }
81  modified = true;
82  outTypes.push_back(shapedTy.cloneWith(shape, shapedTy.getElementType()));
83  }
84 
85  // Return if we don't need to update.
86  if (!modified) {
87  LLVM_DEBUG({ DBGS() << "- op doesn't require update\n"; });
88  return success();
89  }
90 
91  LLVM_DEBUG({
92  DBGS() << "- oldTypes: " << llvm::interleaved_array(op->getResultTypes())
93  << " \n";
94  DBGS() << "- outTypes: " << llvm::interleaved_array(outTypes) << " \n";
95  });
96 
97  // We now have outTypes that need to be turned to cast ops.
98  Location loc = op->getLoc();
99  SmallVector<Value> newResults;
100  // TODO: `mlir::reifyResultShapes` and op verifiers may not agree atm.
101  // This is a confluence problem that will need to be addressed.
102  // For now, we know PadOp and ConcatOp are fine.
103  assert((isa<tensor::PadOp, tensor::ConcatOp>(op.getOperation())) &&
104  "incorrect op");
105  Operation *newOp = rewriter.clone(*op);
106  for (auto [reifiedTy, oldRes] : llvm::zip(outTypes, op->getResults())) {
107  OpResult newRes = newOp->getResult(oldRes.getResultNumber());
108  Type oldTy = oldRes.getType();
109  // Continue if the type remained invariant or is not shaped.
110  if (oldTy == reifiedTy || !isa<MemRefType, RankedTensorType>(oldTy)) {
111  newResults.push_back(newRes);
112  continue;
113  }
114 
115  // Update the type.
116  newRes.setType(reifiedTy);
117  if (isa<RankedTensorType>(reifiedTy)) {
118  newResults.push_back(
119  tensor::CastOp::create(rewriter, loc, oldTy, newRes));
120  } else {
121  assert(isa<MemRefType>(reifiedTy) && "expected a memref type");
122  newResults.push_back(
123  memref::CastOp::create(rewriter, loc, oldTy, newRes));
124  }
125  }
126 
127  LLVM_DEBUG({
128  DBGS() << "- reified results " << llvm::interleaved_array(newResults)
129  << "\n";
130  });
131  rewriter.replaceOp(op, newResults);
132  return success();
133 }
134 
135 //===----------------------------------------------------------------------===//
136 // Pass registration
137 //===----------------------------------------------------------------------===//
138 
139 namespace {
140 struct ReifyResultShapesPass final
141  : public memref::impl::ReifyResultShapesPassBase<ReifyResultShapesPass> {
142  void runOnOperation() override;
143 };
144 } // namespace
145 
146 void ReifyResultShapesPass::runOnOperation() {
148  getOperation()->walk([&](ReifyRankedShapedTypeOpInterface op) {
149  // Handle ops that are not DPS and that do not carry an tied operand shapes.
150  // For now, limit to tensor::PadOp and tensor::ConcatOp.
151  if (!isa<tensor::PadOp, tensor::ConcatOp>(op.getOperation()))
152  return;
153  ops.push_back(op);
154  });
155  IRRewriter rewriter(&getContext());
156  for (ReifyRankedShapedTypeOpInterface op : ops) {
157  rewriter.setInsertionPoint(op);
158  (void)reifyOpResultShapes(rewriter, op);
159  }
160 }
static MLIRContext * getContext(OpFoldResult val)
static LogicalResult reifyOpResultShapes(RewriterBase &rewriter, ReifyRankedShapedTypeOpInterface op)
Reifies the results of op, potentially replacing op with a reified version.
#define DBGS()
This class coordinates rewriting a piece of IR outside of a pattern rewrite, providing a way to keep ...
Definition: PatternMatch.h:764
This class defines the main interface for locations in MLIR and acts as a non-nullable wrapper around...
Definition: Location.h:76
Operation * clone(Operation &op, IRMapping &mapper)
Creates a deep copy of the specified operation, remapping any operands that use values outside of the...
Definition: Builders.cpp:548
This is a value defined by a result of an operation.
Definition: Value.h:447
Operation is the basic unit of execution within MLIR.
Definition: Operation.h:88
OpResult getResult(unsigned idx)
Get the 'idx'th result of this operation.
Definition: Operation.h:407
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...
Instances of the Type class are uniqued, have an immutable identifier and an optional mutable compone...
Definition: Types.h:74
void setType(Type newType)
Mutate the type of this Value to be of the specified type.
Definition: Value.h:116
detail::InFlightRemark failed(Location loc, RemarkOpts opts)
Report an optimization remark that failed.
Definition: Remarks.h:491
Include the generated interface declarations.
std::optional< int64_t > getConstantIntValue(OpFoldResult ofr)
If ofr is a constant integer or an IntegerAttr, return the integer.
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).