161 lines · cpp
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-exception6//7//===----------------------------------------------------------------------===//8//9// This transform reifies result shapes of `ReifyRankedShapedTypeOpInterface`10// operations with ranked `memref` and `tensor` results.11//12//===----------------------------------------------------------------------===//13 14#include "mlir/Dialect/MemRef/Transforms/Passes.h"15 16#include "mlir/Dialect/Affine/IR/AffineOps.h"17#include "mlir/Dialect/MemRef/IR/MemRef.h"18#include "mlir/Dialect/MemRef/Transforms/Transforms.h"19#include "mlir/Dialect/Tensor/IR/Tensor.h"20#include "mlir/Interfaces/InferTypeOpInterface.h"21#include "llvm/Support/InterleavedRange.h"22 23#define DEBUG_TYPE "reify-result-shapes"24#define DBGS() (llvm::dbgs() << "[" DEBUG_TYPE << "]: ")25 26namespace mlir {27namespace memref {28#define GEN_PASS_DEF_REIFYRESULTSHAPESPASS29#include "mlir/Dialect/MemRef/Transforms/Passes.h.inc"30} // namespace memref31} // namespace mlir32 33using namespace mlir;34 35/// Reifies the results of `op`, potentially replacing `op` with a reified36/// version. Returns `failure` if `mlir::reifyResultShapes` returned failure,37/// otherwise it always succeeds. Users of this transform should always expect38/// 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 IR40/// consistent.41static 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 registration137//===----------------------------------------------------------------------===//138 139namespace {140struct ReifyResultShapesPass final141 : public memref::impl::ReifyResultShapesPassBase<ReifyResultShapesPass> {142 void runOnOperation() override;143};144} // namespace145 146void ReifyResultShapesPass::runOnOperation() {147 SmallVector<ReifyRankedShapedTypeOpInterface> ops;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}161