204 lines · cpp
1//===- TypeUtilities.cpp - Helper function for type queries ---------------===//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 defines generic type utilities.10//11//===----------------------------------------------------------------------===//12 13#include "mlir/IR/TypeUtilities.h"14#include "mlir/IR/Attributes.h"15#include "mlir/IR/BuiltinTypes.h"16#include "mlir/IR/Types.h"17#include "mlir/IR/Value.h"18#include "llvm/ADT/SmallVectorExtras.h"19#include <numeric>20 21using namespace mlir;22 23Type mlir::getElementTypeOrSelf(Type type) {24 if (auto st = llvm::dyn_cast<ShapedType>(type))25 return st.getElementType();26 return type;27}28 29Type mlir::getElementTypeOrSelf(Value val) {30 return getElementTypeOrSelf(val.getType());31}32 33Type mlir::getElementTypeOrSelf(Attribute attr) {34 if (auto typedAttr = llvm::dyn_cast<TypedAttr>(attr))35 return getElementTypeOrSelf(typedAttr.getType());36 return {};37}38 39SmallVector<Type, 10> mlir::getFlattenedTypes(TupleType t) {40 SmallVector<Type, 10> fTypes;41 t.getFlattenedTypes(fTypes);42 return fTypes;43}44 45/// Return true if the specified type is an opaque type with the specified46/// dialect and typeData.47bool mlir::isOpaqueTypeWithName(Type type, StringRef dialect,48 StringRef typeData) {49 if (auto opaque = llvm::dyn_cast<mlir::OpaqueType>(type))50 return opaque.getDialectNamespace() == dialect &&51 opaque.getTypeData() == typeData;52 return false;53}54 55/// Returns success if the given two shapes are compatible. That is, they have56/// the same size and each pair of the elements are equal or one of them is57/// dynamic.58LogicalResult mlir::verifyCompatibleShape(ArrayRef<int64_t> shape1,59 ArrayRef<int64_t> shape2) {60 if (shape1.size() != shape2.size())61 return failure();62 for (auto dims : llvm::zip(shape1, shape2)) {63 int64_t dim1 = std::get<0>(dims);64 int64_t dim2 = std::get<1>(dims);65 if (ShapedType::isStatic(dim1) && ShapedType::isStatic(dim2) &&66 dim1 != dim2)67 return failure();68 }69 return success();70}71 72/// Returns success if the given two types have compatible shape. That is,73/// they are both scalars (not shaped), or they are both shaped types and at74/// least one is unranked or they have compatible dimensions. Dimensions are75/// compatible if at least one is dynamic or both are equal. The element type76/// does not matter.77LogicalResult mlir::verifyCompatibleShape(Type type1, Type type2) {78 auto sType1 = llvm::dyn_cast<ShapedType>(type1);79 auto sType2 = llvm::dyn_cast<ShapedType>(type2);80 81 // Either both or neither type should be shaped.82 if (!sType1)83 return success(!sType2);84 if (!sType2)85 return failure();86 87 if (!sType1.hasRank() || !sType2.hasRank())88 return success();89 90 return verifyCompatibleShape(sType1.getShape(), sType2.getShape());91}92 93/// Returns success if the given two arrays have the same number of elements and94/// each pair wise entries have compatible shape.95LogicalResult mlir::verifyCompatibleShapes(TypeRange types1, TypeRange types2) {96 if (types1.size() != types2.size())97 return failure();98 for (auto it : llvm::zip_first(types1, types2))99 if (failed(verifyCompatibleShape(std::get<0>(it), std::get<1>(it))))100 return failure();101 return success();102}103 104LogicalResult mlir::verifyCompatibleDims(ArrayRef<int64_t> dims) {105 if (dims.empty())106 return success();107 auto staticDim =108 llvm::accumulate(dims, dims.front(), [](auto fold, auto dim) {109 return ShapedType::isDynamic(dim) ? fold : dim;110 });111 return success(llvm::all_of(dims, [&](auto dim) {112 return ShapedType::isDynamic(dim) || dim == staticDim;113 }));114}115 116/// Returns success if all given types have compatible shapes. That is, they are117/// all scalars (not shaped), or they are all shaped types and any ranked shapes118/// have compatible dimensions. Dimensions are compatible if all non-dynamic119/// dims are equal. The element type does not matter.120LogicalResult mlir::verifyCompatibleShapes(TypeRange types) {121 auto shapedTypes = llvm::map_to_vector<8>(types, llvm::DynCastTo<ShapedType>);122 // Return failure if some, but not all are not shaped. Return early if none123 // are shaped also.124 if (llvm::none_of(shapedTypes, [](auto t) { return t; }))125 return success();126 if (!llvm::all_of(shapedTypes, [](auto t) { return t; }))127 return failure();128 129 // Return failure if some, but not all, are scalable vectors.130 bool hasScalableVecTypes = false;131 bool hasNonScalableVecTypes = false;132 for (Type t : types) {133 auto vType = llvm::dyn_cast<VectorType>(t);134 if (vType && vType.isScalable())135 hasScalableVecTypes = true;136 else137 hasNonScalableVecTypes = true;138 if (hasScalableVecTypes && hasNonScalableVecTypes)139 return failure();140 }141 142 // Remove all unranked shapes143 auto shapes = llvm::filter_to_vector<8>(144 shapedTypes, [](auto shapedType) { return shapedType.hasRank(); });145 if (shapes.empty())146 return success();147 148 // All ranks should be equal149 auto firstRank = shapes.front().getRank();150 if (llvm::any_of(shapes,151 [&](auto shape) { return firstRank != shape.getRank(); }))152 return failure();153 154 for (unsigned i = 0; i < firstRank; ++i) {155 // Retrieve all ranked dimensions156 auto dims = llvm::map_to_vector<8>(157 llvm::make_filter_range(158 shapes, [&](auto shape) { return shape.getRank() >= i; }),159 [&](auto shape) { return shape.getDimSize(i); });160 if (verifyCompatibleDims(dims).failed())161 return failure();162 }163 164 return success();165}166 167Type OperandElementTypeIterator::mapElement(Value value) const {168 return llvm::cast<ShapedType>(value.getType()).getElementType();169}170 171Type ResultElementTypeIterator::mapElement(Value value) const {172 return llvm::cast<ShapedType>(value.getType()).getElementType();173}174 175TypeRange mlir::insertTypesInto(TypeRange oldTypes, ArrayRef<unsigned> indices,176 TypeRange newTypes,177 SmallVectorImpl<Type> &storage) {178 assert(indices.size() == newTypes.size() &&179 "mismatch between indice and type count");180 if (indices.empty())181 return oldTypes;182 183 auto fromIt = oldTypes.begin();184 for (auto it : llvm::zip(indices, newTypes)) {185 const auto toIt = oldTypes.begin() + std::get<0>(it);186 storage.append(fromIt, toIt);187 storage.push_back(std::get<1>(it));188 fromIt = toIt;189 }190 storage.append(fromIt, oldTypes.end());191 return storage;192}193 194TypeRange mlir::filterTypesOut(TypeRange types, const BitVector &indices,195 SmallVectorImpl<Type> &storage) {196 if (indices.none())197 return types;198 199 for (unsigned i = 0, e = types.size(); i < e; ++i)200 if (!indices[i])201 storage.emplace_back(types[i]);202 return storage;203}204