1 /* Copyright 2022 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_CORE_IR_UTILS_SHAPE_INFERENCE_UTILS_H_ 17 #define TENSORFLOW_CORE_IR_UTILS_SHAPE_INFERENCE_UTILS_H_ 18 19 #include <cstdint> 20 21 #include "mlir/IR/Attributes.h" // from @llvm-project 22 #include "mlir/IR/Operation.h" // from @llvm-project 23 #include "mlir/IR/Types.h" // from @llvm-project 24 #include "mlir/Interfaces/InferTypeOpInterface.h" // from @llvm-project 25 #include "mlir/Support/LLVM.h" // from @llvm-project 26 #include "mlir/Support/LogicalResult.h" // from @llvm-project 27 #include "tensorflow/core/framework/shape_inference.h" 28 29 namespace tensorflow { 30 class OpRegistrationData; 31 } // namespace tensorflow 32 33 namespace mlir { 34 namespace tfg { 35 36 // Function that takes in a value and extracts a constant from it, if available. 37 // If the value cannot be resolved as a constant, a nullptr will be returned. 38 // Certain shape functions require constant values as arguments. 39 using OperandAsConstantFn = llvm::function_ref<Attribute(Value)>; 40 41 // Function that takes in an operation result and computes a shape (can be 42 // partial) value. Certain shape functions require shape values as arguments. 43 using OpResultAsShapeFn = 44 llvm::function_ref<tensorflow::shape_inference::ShapeHandle( 45 tensorflow::shape_inference::InferenceContext&, OpResult)>; 46 47 // Function that takes a result index and returns the element type. Element 48 // types are necessary for handle types (resource, variant). 49 using ResultElementTypeFn = llvm::function_ref<Type(int)>; 50 51 // Extracts the attributes of a MLIR operation and populates the converted 52 // attributes in a proto map<string, AttrValue>. This is used by operation 53 // defined in TF dialect which has different attributes format than TFG dialect. 54 using GetAttrValuesFn = llvm::function_ref<tensorflow::Status( 55 Operation*, llvm::StringRef, const tensorflow::OpRegistrationData*, bool, 56 tensorflow::AttrValueMap*)>; 57 58 // Runs TensorFlow shape inference associated to the op type registered in the 59 // TensorFlow op registry based on the Graph version, operands, and attributes. 60 // Invoking this shape function will create conversions of parameters to the 61 // TensorFlow Graph equivalent data structures and back to MLIR equivalent data 62 // structures. This does not use a natively implemented shape inference in MLIR, 63 // and instead is temporary until shape functions are reimplemented/migrated to 64 // being in MLIR instead of the TensorFlow op registry. 65 // Note that the default way to get the attrs in the operation is using the API 66 // in TFG importer. For operations that has different format of attributes, they 67 // should give the `get_attr_values_fn` to read the attributes correctly. 68 LogicalResult InferReturnTypeComponentsForTFOp( 69 Optional<Location> location, Operation* op, ValueRange operands, 70 int64_t graph_version, OperandAsConstantFn operand_as_constant_fn, 71 OpResultAsShapeFn op_result_as_shape_fn, 72 ResultElementTypeFn result_element_type_fn, 73 GetAttrValuesFn get_attr_values_fn, 74 SmallVectorImpl<ShapedTypeComponents>& inferred_return_shapes); 75 76 // This one is almost the same as the above one, the difference is that we use 77 // ConvertOperationToNode to convert the operation to NodeDef to get the attr 78 // values. 79 LogicalResult InferReturnTypeComponentsForTFOp( 80 Optional<Location> location, Operation* op, ValueRange operands, 81 int64_t graph_version, OperandAsConstantFn operand_as_constant_fn, 82 OpResultAsShapeFn op_result_as_shape_fn, 83 ResultElementTypeFn result_element_type_fn, 84 SmallVectorImpl<ShapedTypeComponents>& inferred_return_shapes); 85 86 } // namespace tfg 87 } // namespace mlir 88 89 #endif // TENSORFLOW_CORE_IR_UTILS_SHAPE_INFERENCE_UTILS_H_ 90