1 /* Copyright 2020 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 #ifndef MLIR_HLO_UTILS_BROADCAST_UTILS_H 17 #define MLIR_HLO_UTILS_BROADCAST_UTILS_H 18 19 // Utilities relating to implementing HLO broadcasting. 20 // Note: This file should not depend on any non-MLIR TensorFlow libraries. 21 22 #include "mlir/IR/Attributes.h" 23 #include "mlir/IR/Builders.h" 24 #include "mlir/IR/BuiltinTypes.h" 25 #include "mlir/IR/Location.h" 26 #include "mlir/Interfaces/InferTypeOpInterface.h" 27 #include "mlir/Support/LLVM.h" 28 29 namespace mlir { 30 namespace hlo { 31 32 // Checks whether the given operand types and broadcast_dims attr represent a 33 // legal combination for "numpy" style broadcasting (where 1-dims are prepended 34 // to the smaller ranked operand until it is of the same rank as the larger). 35 // See: https://docs.scipy.org/doc/numpy/reference/ufuncs.html 36 bool isLegalNumpyRankedBroadcast(Value lhs, Value rhs, 37 DenseIntElementsAttr broadcastDims); 38 39 // Emits shape dialect ops to compute the result shape for a broadcasting 40 // binary/n-ary elementwise op which broadcasts according to "numpy" semantics 41 // (see above), returning an extent tensor of the resulting shape. The function 42 // should only be used in contexts that ensure both operands to be 43 // broadcastable. 44 Value computeBinaryElementwiseBroadcastingResultExtents(Location loc, Value lhs, 45 Value rhs, 46 OpBuilder& builder); 47 Value computeNaryElementwiseBroadcastingResultExtents(Location loc, 48 ValueRange operands, 49 OpBuilder& builder); 50 51 } // namespace hlo 52 } // namespace mlir 53 54 #endif // MLIR_HLO_UTILS_BROADCAST_UTILS_H 55