1 /* Copyright 2019 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_TENSORFLOW_TRANSFORMS_BRIDGE_H_ 17 #define TENSORFLOW_COMPILER_MLIR_TENSORFLOW_TRANSFORMS_BRIDGE_H_ 18 19 #include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project 20 #include "mlir/IR/BuiltinOps.h" // from @llvm-project 21 #include "tensorflow/core/lib/core/status.h" 22 23 namespace mlir { 24 namespace TFTPU { 25 26 // Run all the passes involved in transforming the graph before execution so 27 // that it is suitable for targeting TPUs. When enable_logging is true, enables 28 // tensorflow::BridgeLogger. When fallback_enabled is true, it means if the 29 // bridge fails the old bridge will run. This is used for logging and doesn't 30 // affect any logic. 31 tensorflow::Status TPUBridge(ModuleOp module, bool enable_logging, 32 bool fallback_enabled = false); 33 34 // Run all the passes involved in transforming the graph before execution so 35 // that it is suitable for targeting TPUs. When enable_logging is true, enables 36 // tensorflow::BridgeLogger. When fallback_enabled is true, it means if the 37 // bridge fails the old bridge will run. This is used for logging and doesn't 38 // affect any logic. 39 // This variant of `TPUBridge` is intended for TensorFlow V1 compatibility. 40 tensorflow::Status TPUBridgeV1Compat(ModuleOp module, bool enable_logging, 41 bool fallback_enabled = false); 42 43 } // namespace TFTPU 44 45 namespace TF { 46 47 // Runs all passes involved in transforming or optimizing an MLIR graph without 48 // any target specialization. When enable_logging is true, enables 49 // tensorflow::BridgeLogger. When enable_inliner is true, enables the inliner 50 // pass. 51 tensorflow::Status RunBridgeWithStandardPipeline(ModuleOp module, 52 bool enable_logging, 53 bool enable_inliner); 54 55 // Runs all passes for non TPU (GPU and CPU) graph. 56 tensorflow::Status RunTFXLABridge(ModuleOp module, bool enable_logging); 57 } // namespace TF 58 59 } // namespace mlir 60 61 #endif // TENSORFLOW_COMPILER_MLIR_TENSORFLOW_TRANSFORMS_BRIDGE_H_ 62