1 /* Copyright 2021 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_TF_OP_WRAPPER_H_ 17 #define TENSORFLOW_CORE_IR_TF_OP_WRAPPER_H_ 18 19 #include "llvm/ADT/iterator_range.h" 20 #include "mlir/IR/Operation.h" // from @llvm-project 21 #include "mlir/IR/OperationSupport.h" // from @llvm-project 22 #include "mlir/IR/TypeRange.h" // from @llvm-project 23 #include "tensorflow/core/ir/dialect.h" 24 #include "tensorflow/core/ir/ops.h" 25 26 namespace mlir { 27 namespace tfg { 28 29 // Wrapper class exposing convenience methods to manipulate TensorFlow graph 30 // nodes uniformly. 31 class TFOp { 32 public: 33 explicit TFOp(Operation &op); 34 35 // Returns a pointer to the TensorFlow Graph Dialect. It nevers returns 36 // nullptr. getDialect()37 TFGraphDialect *getDialect() { 38 return static_cast<TFGraphDialect *>(op_.getDialect()); 39 } 40 41 // The control operands are always after the regular inputs. getControlOperands()42 ValueRange getControlOperands() { 43 OperandRange operands = op_.getOperands(); 44 if (operands.empty()) return operands; 45 OperandRange::iterator first_control_op = std::prev(operands.end()); 46 while (first_control_op != operands.begin() && 47 first_control_op.getBase()->get().getType().isa<ControlType>()) 48 --first_control_op; 49 if (!first_control_op.getBase()->get().getType().isa<ControlType>()) 50 ++first_control_op; 51 return llvm::make_range(first_control_op, operands.end()); 52 } 53 54 // Returns the control token produced by this operation. controlRet()55 Value controlRet() { return op_.getResult(op_.getNumResults() - 1); } 56 57 // Returns the node name for this operation. 58 StringAttr nameAttr(); 59 StringRef name(); 60 61 // Returns the requested device, which is also the "device" field in a 62 // GraphDef. 63 StringAttr requestedDeviceAttr(); 64 StringRef requestedDevice(); 65 66 // Returns the assigned device, this field is set by placer in general. 67 StringAttr assignedDeviceAttr(); 68 StringRef assignedDevice(); 69 70 // Returns the device, preferring the assigned device if set, and the 71 // requested device otherwise. deviceAttr()72 StringAttr deviceAttr() { 73 StringAttr device = assignedDeviceAttr(); 74 if (device) return device; 75 return requestedDeviceAttr(); 76 } device()77 StringRef device() { 78 StringAttr device_attr = deviceAttr(); 79 if (device_attr) return device_attr.getValue(); 80 return ""; 81 } 82 83 // Forward `->` to the underlying operation, exposing the `Operation` methods. 84 Operation *operator->() { return &op_; } 85 Operation &operator*() { return op_; } 86 87 private: 88 Operation &op_; 89 }; 90 91 // Return the function called by this operation if this operation is a call and 92 // the function exists in the symbol_table, otherwise return null. 93 GraphFuncOp getCalledFunction(Operation *op, SymbolTable &symbol_table); 94 95 } // namespace tfg 96 } // namespace mlir 97 98 #endif // TENSORFLOW_CORE_IR_TF_OP_WRAPPER_H_ 99