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 TENSORFLOW_COMPILER_MLIR_HLO_INCLUDE_MLIR_HLO_UTILS_BROADCAST_UTILS_H_ 17 #define TENSORFLOW_COMPILER_MLIR_HLO_INCLUDE_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 broadcast_dims); 38 39 // Emits shape dialect ops to compute the result shape for a broadcasting 40 // binary elementwise op which broadcasts according to "numpy" semantics 41 // (see above), returning a `shape.shape` or an extent tensor of the resulting 42 // shape. The result should only be an extent tensor in contexts that ensure 43 // both operands to be broadcastable. 44 Value ComputeBinaryElementwiseBroadcastingResultExtents( 45 Location loc, Value lhs, Value rhs, OpBuilder& builder, 46 bool unsafe_as_extent_tensor); 47 48 } // namespace hlo 49 } // namespace mlir 50 51 #endif // TENSORFLOW_COMPILER_MLIR_HLO_INCLUDE_MLIR_HLO_UTILS_BROADCAST_UTILS_H_ 52