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_TRANSFORMS_UTILS_UTILS_H_ 17 #define TENSORFLOW_CORE_TRANSFORMS_UTILS_UTILS_H_ 18 19 #include "mlir/IR/Types.h" // from @llvm-project 20 #include "mlir/IR/Value.h" // from @llvm-project 21 #include "tensorflow/core/ir/dialect.h" 22 23 namespace mlir { 24 25 class Operation; 26 class NamedAttrList; 27 28 namespace tfg { 29 namespace util { 30 31 // Returns true if the op has the requested device attribute. 32 bool OpHasDevice(Operation *op, const char *device_name); 33 34 // Erase the attribute starts with "_". 35 void EraseRegularNodeAttributes(NamedAttrList &attr_list); 36 37 // When rewriting an operation 1-to-1, intrinsic attributes are manually 38 // forwarded, modified, or dropped. For example, when `If` is rewritten to 39 // `IfRegion`, 40 // 41 // 1. `Tout` is forwarded as is, 42 // 2. `then_branch` is changed to `then_attrs` which contain the attribute 43 // dictionary part of the `#tf_type.func`, and 44 // 3. `Tin` is dropped. 45 // 46 // Non-intrinsic attributes, e.g. `_tpu_cluster`, are blindly forwarded to the 47 // new operation. 48 void ForwardNonIntrinsicAttributes(Operation *src, Operation *dst); 49 50 // Add an argument to a loop region. This inserts the new data argument and 51 // control argument at the correct positions and returns them. Also, this 52 // function updates any preserved argument attributes by inserting a null. 53 struct LoopRegionArgumentUpdate { 54 BlockArgument data, ctl; 55 }; 56 LoopRegionArgumentUpdate LoopRegionAddArgument(Region ®ion, Type type); 57 58 // Erase an argument from a loop region. This erases the corresponding control 59 // argument. Also, this function updates any preserved argument attributes by 60 // deleting them. 61 void LoopRegionEraseArgument(Region ®ion, unsigned index); 62 63 // Indicate that a result has been added to a loop region. Call this function to 64 // update the preserved result attributes. 65 void LoopRegionResultAdded(Region ®ion, unsigned num = 1); 66 67 // Indicate that a result has been erased from a loop region. Call this function 68 // to update the preserved result attributes. 69 void LoopRegionResultErased(Region ®ion, unsigned index); 70 71 // Erase operands from an op that might have an `operand_segment_sizes` , 72 // updating the attribute in-place if present. 73 void SizedOperandSegmentsEraseOperands(Operation *op, 74 ArrayRef<unsigned> indices); 75 void SizedOperandSegmentsEraseOperands(Operation *op, 76 const llvm::BitVector &erase); 77 78 } // namespace util 79 } // namespace tfg 80 } // namespace mlir 81 82 #endif // TENSORFLOW_CORE_TRANSFORMS_UTILS_UTILS_H_ 83