271 lines · cpp
1//===- InferTypeOpInterface.cpp - Infer Type Interfaces ---------*- 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-exception6//7//===----------------------------------------------------------------------===//8//9// This file contains the definitions of the infer op interfaces defined in10// `InferTypeOpInterface.td`.11//12//===----------------------------------------------------------------------===//13 14#include "mlir/Interfaces/InferTypeOpInterface.h"15#include "mlir/IR/BuiltinTypes.h"16#include "mlir/IR/Matchers.h"17#include "llvm/Support/FormatVariadic.h"18#include "llvm/Support/InterleavedRange.h"19 20using namespace mlir;21 22namespace mlir {23#include "mlir/Interfaces/InferTypeOpInterface.cpp.inc"24} // namespace mlir25 26LogicalResult27mlir::reifyResultShapes(OpBuilder &b, Operation *op,28 ReifiedRankedShapedTypeDims &reifiedReturnShapes) {29 auto reifiableOp = dyn_cast<ReifyRankedShapedTypeOpInterface>(op);30 if (!reifiableOp)31 return failure();32 LogicalResult status = reifiableOp.reifyResultShapes(b, reifiedReturnShapes);33#ifndef NDEBUG34 if (failed(status))35 return failure();36 // Assert that ReifyRankedShapedTypeOpInterface::reifyResultShapes produced37 // a correct result.38 int64_t resultIdx = 0;39 for (OpResult result : op->getResults()) {40 auto shapedType = dyn_cast<ShapedType>(result.getType());41 if (!shapedType)42 continue;43 if (!shapedType.hasRank()) {44 // Nothing to check for unranked shaped values.45 ++resultIdx;46 continue;47 }48 // Assert one OpFoldResult per dimension.49 assert(shapedType.getRank() ==50 static_cast<int64_t>(reifiedReturnShapes[resultIdx].size()) &&51 "incorrect implementation of ReifyRankedShapedTypeOpInterface");52 ++resultIdx;53 }54 // Assert that every shaped value result was reified.55 assert(resultIdx == static_cast<int64_t>(reifiedReturnShapes.size()) &&56 "incorrect implementation of ReifyRankedShapedTypeOpInterface");57#endif // NDEBUG58 return status;59}60 61FailureOr<SmallVector<OpFoldResult>>62mlir::reifyShapeOfResult(OpBuilder &b, Operation *op, int resultIndex) {63 auto reifiableOp = dyn_cast<ReifyRankedShapedTypeOpInterface>(op);64 if (!reifiableOp)65 return failure();66 return reifiableOp.reifyShapeOfResult(b, resultIndex);67}68 69FailureOr<OpFoldResult> mlir::reifyDimOfResult(OpBuilder &b, Operation *op,70 int resultIndex, int dim) {71 auto reifiableOp = dyn_cast<ReifyRankedShapedTypeOpInterface>(op);72 if (!reifiableOp)73 return failure();74 return reifiableOp.reifyDimOfResult(b, resultIndex, dim);75}76 77bool ShapeAdaptor::hasRank() const {78 if (val.isNull())79 return false;80 if (auto t = llvm::dyn_cast_if_present<Type>(val))81 return cast<ShapedType>(t).hasRank();82 if (isa<Attribute>(val))83 return true;84 return cast<ShapedTypeComponents *>(val)->hasRank();85}86 87Type ShapeAdaptor::getElementType() const {88 if (val.isNull())89 return nullptr;90 if (auto t = llvm::dyn_cast_if_present<Type>(val))91 return cast<ShapedType>(t).getElementType();92 if (isa<Attribute>(val))93 return nullptr;94 return cast<ShapedTypeComponents *>(val)->getElementType();95}96 97void ShapeAdaptor::getDims(SmallVectorImpl<int64_t> &res) const {98 assert(hasRank());99 if (auto t = llvm::dyn_cast_if_present<Type>(val)) {100 ArrayRef<int64_t> vals = cast<ShapedType>(t).getShape();101 res.assign(vals.begin(), vals.end());102 } else if (auto attr = llvm::dyn_cast_if_present<Attribute>(val)) {103 auto dattr = cast<DenseIntElementsAttr>(attr);104 res.clear();105 res.reserve(dattr.size());106 for (auto it : dattr.getValues<APInt>())107 res.push_back(it.getSExtValue());108 } else {109 auto vals = cast<ShapedTypeComponents *>(val)->getDims();110 res.assign(vals.begin(), vals.end());111 }112}113 114void ShapeAdaptor::getDims(ShapedTypeComponents &res) const {115 assert(hasRank());116 res.ranked = true;117 getDims(res.dims);118}119 120int64_t ShapeAdaptor::getDimSize(int index) const {121 assert(hasRank());122 if (auto t = llvm::dyn_cast_if_present<Type>(val))123 return cast<ShapedType>(t).getDimSize(index);124 if (auto attr = llvm::dyn_cast_if_present<Attribute>(val))125 return cast<DenseIntElementsAttr>(attr)126 .getValues<APInt>()[index]127 .getSExtValue();128 auto *stc = cast<ShapedTypeComponents *>(val);129 return stc->getDims()[index];130}131 132int64_t ShapeAdaptor::getRank() const {133 assert(hasRank());134 if (auto t = llvm::dyn_cast_if_present<Type>(val))135 return cast<ShapedType>(t).getRank();136 if (auto attr = llvm::dyn_cast_if_present<Attribute>(val))137 return cast<DenseIntElementsAttr>(attr).size();138 return cast<ShapedTypeComponents *>(val)->getDims().size();139}140 141bool ShapeAdaptor::hasStaticShape() const {142 if (!hasRank())143 return false;144 145 if (auto t = llvm::dyn_cast_if_present<Type>(val))146 return cast<ShapedType>(t).hasStaticShape();147 if (auto attr = llvm::dyn_cast_if_present<Attribute>(val)) {148 auto dattr = cast<DenseIntElementsAttr>(attr);149 for (auto index : dattr.getValues<APInt>())150 if (ShapedType::isDynamic(index.getSExtValue()))151 return false;152 return true;153 }154 auto *stc = cast<ShapedTypeComponents *>(val);155 return llvm::none_of(stc->getDims(), ShapedType::isDynamic);156}157 158int64_t ShapeAdaptor::getNumElements() const {159 assert(hasStaticShape() && "cannot get element count of dynamic shaped type");160 161 if (auto t = llvm::dyn_cast_if_present<Type>(val))162 return cast<ShapedType>(t).getNumElements();163 164 if (auto attr = llvm::dyn_cast_if_present<Attribute>(val)) {165 auto dattr = cast<DenseIntElementsAttr>(attr);166 int64_t num = 1;167 for (auto index : dattr.getValues<APInt>()) {168 num *= index.getZExtValue();169 assert(num >= 0 && "integer overflow in element count computation");170 }171 return num;172 }173 174 auto *stc = cast<ShapedTypeComponents *>(val);175 int64_t num = 1;176 for (int64_t dim : stc->getDims()) {177 num *= dim;178 assert(num >= 0 && "integer overflow in element count computation");179 }180 return num;181}182 183void ShapeAdaptor::dump() const {184 if (!hasRank()) {185 llvm::errs() << "<<unranked>>\n";186 return;187 }188 189 SmallVector<int64_t> dims;190 getDims(dims);191 auto mapped = llvm::map_range(dims, [](int64_t dim) -> std::string {192 if (ShapedType::isDynamic(dim))193 return "?";194 return llvm::formatv("{0}", dim).str();195 });196 llvm::errs() << "rank = " << getRank()197 << " dims = " << llvm::interleaved_array(mapped, "x") << "\n";198}199 200ShapeAdaptor ValueShapeRange::getValueAsShape(int index) {201 Value val = operator[](index);202 if (valueToShape)203 if (ShapeAdaptor ret = valueToShape(val))204 return ret;205 206 DenseIntElementsAttr attr;207 if (!matchPattern(val, m_Constant(&attr)))208 return nullptr;209 if (attr.getType().getRank() != 1)210 return nullptr;211 return attr;212}213 214ShapeAdaptor ValueShapeRange::getShape(Value val) const {215 if (operandShape)216 if (ShapeAdaptor ret = operandShape(val))217 return ret;218 return val.getType();219}220 221ShapeAdaptor ValueShapeRange::getShape(int index) const {222 if (index < 0 || static_cast<size_t>(index) >= size())223 return nullptr;224 return getShape(operator[](index));225}226 227LogicalResult mlir::detail::inferReturnTensorTypes(228 ArrayRef<ShapedTypeComponents> retComponents,229 SmallVectorImpl<Type> &inferredReturnTypes) {230 for (const auto &shapeAndType : retComponents) {231 Type elementTy = shapeAndType.getElementType();232 assert(elementTy && "element type required to construct tensor");233 234 Attribute attr = shapeAndType.getAttribute();235 if (shapeAndType.hasRank()) {236 inferredReturnTypes.push_back(237 RankedTensorType::get(shapeAndType.getDims(), elementTy, attr));238 } else {239 assert(attr == nullptr && "attribute not supported");240 inferredReturnTypes.push_back(UnrankedTensorType::get(elementTy));241 }242 }243 return success();244}245 246LogicalResult mlir::detail::verifyInferredResultTypes(Operation *op) {247 SmallVector<Type, 4> inferredReturnTypes(op->getResultTypes());248 auto retTypeFn = cast<InferTypeOpInterface>(op);249 auto result = retTypeFn.refineReturnTypes(250 op->getContext(), op->getLoc(), op->getOperands(),251 op->getRawDictionaryAttrs(), op->getPropertiesStorage(), op->getRegions(),252 inferredReturnTypes);253 if (failed(result))254 op->emitOpError() << "failed to infer returned types";255 256 return result;257}258 259void mlir::detail::reportFatalInferReturnTypesError(OperationState &state) {260 std::string buffer;261 llvm::raw_string_ostream os(buffer);262 os << "Failed to infer result type(s):\n"263 << "\"" << state.name << "\"(...) "264 << state.attributes.getDictionary(state.location.getContext()) << " : ("265 << llvm::interleaved(llvm::map_range(266 state.operands, [](Value val) { return val.getType(); }))267 << ") -> ( ??? )";268 emitRemark(state.location, "location of op");269 llvm::report_fatal_error(llvm::StringRef(buffer));270}271