1 /* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
2
3 Licensed under the Apache License, Version 2.0 (the "License");
4 you may not use this file except in compliance with the License.
5 You may obtain a copy of the License at
6
7 http://www.apache.org/licenses/LICENSE-2.0
8
9 Unless required by applicable law or agreed to in writing, software
10 distributed under the License is distributed on an "AS IS" BASIS,
11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 See the License for the specific language governing permissions and
13 limitations under the License.
14 ==============================================================================*/
15
16 #include "tensorflow/compiler/mlir/xla/ir/hlo_utils.h"
17
18 #include <numeric>
19
20 #include "mlir/IR/Attributes.h" // TF:llvm-project
21
22 namespace mlir {
23 namespace xla {
24
getBroadcastDimensionsAttr(Builder * b,Value x,Value y)25 DenseIntElementsAttr getBroadcastDimensionsAttr(Builder *b, Value x, Value y) {
26 TensorType xType = x.getType().dyn_cast<RankedTensorType>();
27 TensorType yType = y.getType().dyn_cast<RankedTensorType>();
28 if (xType == yType || !xType || !yType) return {};
29
30 // If the shapes have the same rank, then there is nothing to do.
31 auto xRank = xType.getRank(), yRank = yType.getRank();
32 if (xRank == yRank) return {};
33
34 // Otherwise if the ranks of the inputs don't match, TensorFlow automatically
35 // reshapes the smaller by padding with dimensions of size 1 as a prefix. In
36 // other words to pad a 5-vector to a 3-dimensional tensor it is reshaped to
37 // have shape [1,1,5]. XLA's automatic broadcast code is able to broadcast
38 // from lower to higher rank, but doesn't assume you want to pad as a prefix
39 // of the dimensions, and instead needs to be told which dimensions of the
40 // higher rank tensor to match to the lower rank tensor.
41 auto maxRank = std::max(xRank, yRank);
42 auto minRank = std::min(xRank, yRank);
43
44 // Match the lower rank tensor along the larger-numbered dimensions of the
45 // higher rank tensor.
46 SmallVector<int64_t, 4> broadcastDimensions(minRank);
47 std::iota(broadcastDimensions.begin(), broadcastDimensions.end(),
48 maxRank - minRank);
49
50 RankedTensorType type =
51 RankedTensorType::get({minRank}, b->getIntegerType(64));
52 return DenseIntElementsAttr::get(type, broadcastDimensions);
53 }
54
GetScalarOfType(Type ty,int64_t raw_value)55 DenseElementsAttr GetScalarOfType(Type ty, int64_t raw_value) {
56 RankedTensorType scalar_ty = RankedTensorType::get({}, ty);
57
58 if (auto float_ty = ty.dyn_cast<FloatType>()) {
59 APFloat value(float_ty.getFloatSemantics(), raw_value);
60 return DenseElementsAttr::get(scalar_ty, value);
61 }
62 auto int_ty = ty.cast<IntegerType>();
63 APInt value(int_ty.getWidth(), static_cast<int64_t>(raw_value), true);
64 return DenseElementsAttr::get(scalar_ty, value);
65 }
66
67 } // namespace xla
68 } // namespace mlir
69