1//===- Traits.cpp - Common op traits shared by dialects -------------------===//
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#include "mlir/Dialect/Traits.h"
10#include "mlir/IR/BuiltinTypes.h"
11#include "mlir/IR/TypeUtilities.h"
12#include <optional>
13
14using namespace mlir;
15
16bool OpTrait::util::staticallyKnownBroadcastable(ArrayRef<int64_t> shape1,
17 ArrayRef<int64_t> shape2) {
18 SmallVector<SmallVector<int64_t, 6>, 2> extents;
19 extents.emplace_back(Args: shape1.begin(), Args: shape1.end());
20 extents.emplace_back(Args: shape2.begin(), Args: shape2.end());
21 return staticallyKnownBroadcastable(shapes: extents);
22}
23
24bool OpTrait::util::staticallyKnownBroadcastable(
25 ArrayRef<SmallVector<int64_t, 6>> shapes) {
26 assert(!shapes.empty() && "Expected at least one shape");
27 size_t maxRank = shapes[0].size();
28 for (size_t i = 1; i != shapes.size(); ++i)
29 maxRank = std::max(a: maxRank, b: shapes[i].size());
30
31 // We look backwards through every column of `shapes`.
32 for (size_t i = 0; i != maxRank; ++i) {
33 bool seenDynamic = false;
34 std::optional<int64_t> nonOneDim;
35 for (ArrayRef<int64_t> extent : shapes) {
36 int64_t dim = i >= extent.size() ? 1 : extent[extent.size() - i - 1];
37
38 if (dim == 1)
39 continue;
40
41 // Dimensions are compatible when
42 //. 1. One is dynamic, the rest are 1
43 if (ShapedType::isDynamic(dValue: dim)) {
44 if (seenDynamic || nonOneDim)
45 return false;
46 seenDynamic = true;
47 }
48
49 // 2. All are 1 or a specific constant.
50 if (nonOneDim && dim != *nonOneDim)
51 return false;
52
53 nonOneDim = dim;
54 }
55 }
56 return true;
57}
58
59bool OpTrait::util::getBroadcastedShape(ArrayRef<int64_t> shape1,
60 ArrayRef<int64_t> shape2,
61 SmallVectorImpl<int64_t> &resultShape) {
62 // To compute the result broadcasted shape, we compare operand shapes
63 // element-wise: starting with the trailing dimensions, and working the
64 // way backward. Two dimensions are compatible when
65 // 1. they are equal, or
66 // 2. one of them is 1
67 // The result shape has the maximum among the two inputs at every
68 // dimension index.
69
70 resultShape.clear();
71 if (shape1.size() > shape2.size()) {
72 llvm::append_range(C&: resultShape, R&: shape1);
73 } else {
74 llvm::append_range(C&: resultShape, R&: shape2);
75 }
76
77 auto i1 = shape1.rbegin(), e1 = shape1.rend();
78 auto i2 = shape2.rbegin(), e2 = shape2.rend();
79 auto iR = resultShape.rbegin();
80
81 // Check each dimension is consistent.
82 for (; i1 != e1 && i2 != e2; ++i1, ++i2, ++iR) {
83 if (ShapedType::isDynamic(dValue: *i1) || ShapedType::isDynamic(dValue: *i2)) {
84 // One or both dimensions is unknown. Follow TensorFlow behavior:
85 // - If either dimension is greater than 1, we assume that the program is
86 // correct, and the other dimension will be broadcasted to match it.
87 // - If either dimension is 1, the other dimension is the output.
88 if (*i1 > 1) {
89 *iR = *i1;
90 } else if (*i2 > 1) {
91 *iR = *i2;
92 } else if (*i1 == 1) {
93 *iR = *i2;
94 } else if (*i2 == 1) {
95 *iR = *i1;
96 } else {
97 *iR = ShapedType::kDynamic;
98 }
99 } else {
100 if (*i1 == *i2 || *i2 == 1) {
101 *iR = *i1;
102 } else if (*i1 == 1) {
103 *iR = *i2;
104 } else {
105 // This dimension of the two operand types is incompatible.
106 resultShape.clear();
107 return false;
108 }
109 }
110 }
111
112 return true;
113}
114
115/// Returns the shape of the given type. Scalars will be considered as having a
116/// shape with zero dimensions.
117static ArrayRef<int64_t> getShape(Type type) {
118 if (auto sType = dyn_cast<ShapedType>(Val&: type))
119 return sType.getShape();
120 return {};
121}
122
123/// Returns the result broadcast composition type from the two given types by
124/// following NumPy broadcast semantics. Returned type may have dynamic shape if
125/// either of the input types has dynamic shape. Returns null type if the two
126/// given types are not broadcast-compatible.
127///
128/// elementType, if specified, will be used as the element type of the
129/// broadcasted result type. Otherwise it is required that the element type of
130/// type1 and type2 is the same and this element type will be used as the
131/// resultant element type.
132Type OpTrait::util::getBroadcastedType(Type type1, Type type2,
133 Type elementType) {
134 // If the elementType is not specified, then the use the common element type
135 // of the inputs or fail if there is no common element type.
136 if (!elementType) {
137 elementType = getElementTypeOrSelf(type: type1);
138 if (elementType != getElementTypeOrSelf(type: type2))
139 return {};
140 }
141
142 // If one of the types is unranked tensor, then the other type shouldn't be
143 // vector and the result should have unranked tensor type.
144 if (isa<UnrankedTensorType>(Val: type1) || isa<UnrankedTensorType>(Val: type2)) {
145 if (isa<VectorType>(Val: type1) || isa<VectorType>(Val: type2))
146 return {};
147 return UnrankedTensorType::get(elementType);
148 }
149
150 // Returns the type kind if the given type is a vector or ranked tensor type.
151 // Returns std::nullopt otherwise.
152 auto getCompositeTypeKind = [](Type type) -> std::optional<TypeID> {
153 if (isa<VectorType, RankedTensorType>(Val: type))
154 return type.getTypeID();
155 return std::nullopt;
156 };
157
158 // Make sure the composite type, if has, is consistent.
159 std::optional<TypeID> compositeKind1 = getCompositeTypeKind(type1);
160 std::optional<TypeID> compositeKind2 = getCompositeTypeKind(type2);
161 std::optional<TypeID> resultCompositeKind;
162
163 if (compositeKind1 && compositeKind2) {
164 // Disallow mixing vector and tensor.
165 if (compositeKind1 != compositeKind2)
166 return {};
167 resultCompositeKind = compositeKind1;
168 } else if (compositeKind1) {
169 resultCompositeKind = compositeKind1;
170 } else if (compositeKind2) {
171 resultCompositeKind = compositeKind2;
172 }
173
174 // Get the shape of each type.
175 SmallVector<int64_t, 4> resultShape;
176 if (!getBroadcastedShape(shape1: getShape(type: type1), shape2: getShape(type: type2), resultShape))
177 return {};
178
179 // Compose the final broadcasted type
180 if (resultCompositeKind == VectorType::getTypeID())
181 return VectorType::get(shape: resultShape, elementType);
182 if (resultCompositeKind == RankedTensorType::getTypeID())
183 return RankedTensorType::get(shape: resultShape, elementType);
184 return elementType;
185}
186
187/// Returns a tuple corresponding to whether range has tensor or vector type.
188template <typename iterator_range>
189static std::tuple<bool, bool> hasTensorOrVectorType(iterator_range types) {
190 return {llvm::any_of(types, llvm::IsaPred<TensorType>),
191 llvm::any_of(types, llvm::IsaPred<VectorType>)};
192}
193
194static bool isCompatibleInferredReturnShape(ArrayRef<int64_t> inferred,
195 ArrayRef<int64_t> existing) {
196 // If both interred and existing dimensions are static, they must be equal.
197 auto isCompatible = [](int64_t inferredDim, int64_t existingDim) {
198 return ShapedType::isDynamic(dValue: existingDim) ||
199 ShapedType::isDynamic(dValue: inferredDim) || inferredDim == existingDim;
200 };
201 if (inferred.size() != existing.size())
202 return false;
203 for (auto [inferredDim, existingDim] : llvm::zip_equal(t&: inferred, u&: existing))
204 if (!isCompatible(inferredDim, existingDim))
205 return false;
206 return true;
207}
208
209static std::string getShapeString(ArrayRef<int64_t> shape) {
210 // TODO: should replace with printing shape more uniformly across here and
211 // when in type.
212 std::string ret;
213 llvm::raw_string_ostream ss(ret);
214 ss << '\'';
215 llvm::interleave(
216 c: shape, os&: ss,
217 each_fn: [&](int64_t dim) {
218 if (ShapedType::isDynamic(dValue: dim))
219 ss << '?';
220 else
221 ss << dim;
222 },
223 separator: "x");
224 ss << '\'';
225 return ret;
226}
227
228LogicalResult OpTrait::impl::verifyCompatibleOperandBroadcast(Operation *op) {
229 // Ensure broadcasting only tensor or only vector types.
230 auto operandsHasTensorVectorType =
231 hasTensorOrVectorType(types: op->getOperandTypes());
232 auto resultsHasTensorVectorType = hasTensorOrVectorType(types: op->getResultTypes());
233 if ((std::get<0>(t&: operandsHasTensorVectorType) ||
234 std::get<0>(t&: resultsHasTensorVectorType)) &&
235 (std::get<1>(t&: operandsHasTensorVectorType) ||
236 std::get<1>(t&: resultsHasTensorVectorType)))
237 return op->emitError(message: "cannot broadcast vector with tensor");
238
239 auto rankedOperands =
240 make_filter_range(Range: op->getOperandTypes(), Pred: llvm::IsaPred<RankedTensorType>);
241
242 // If all operands are unranked, then all result shapes are possible.
243 if (rankedOperands.empty())
244 return success();
245
246 // Compute broadcasted shape of operands (which requires that operands are
247 // broadcast compatible). The results need to be broadcast compatible with
248 // this result shape.
249 SmallVector<int64_t, 4> resultShape;
250 (void)util::getBroadcastedShape(shape1: getShape(type: *rankedOperands.begin()), shape2: {},
251 resultShape);
252 for (auto other : make_early_inc_range(Range&: rankedOperands)) {
253 SmallVector<int64_t, 4> temp = resultShape;
254 if (!util::getBroadcastedShape(shape1: temp, shape2: getShape(type: other), resultShape))
255 return op->emitOpError(message: "operands don't have broadcast-compatible shapes");
256 }
257
258 auto rankedResults =
259 make_filter_range(Range: op->getResultTypes(), Pred: llvm::IsaPred<RankedTensorType>);
260
261 // If all of the results are unranked then no further verification.
262 if (rankedResults.empty())
263 return success();
264
265 for (auto type : rankedResults) {
266 ArrayRef<int64_t> actualSuffix =
267 getShape(type).take_back(N: resultShape.size());
268 if (!isCompatibleInferredReturnShape(inferred: resultShape, existing: actualSuffix))
269 return op->emitOpError()
270 << "result type " << getShapeString(shape: getShape(type))
271 << " not broadcast compatible with broadcasted operands's shapes "
272 << getShapeString(shape: resultShape);
273 }
274 return success();
275}
276

source code of mlir/lib/Dialect/Traits.cpp