| 36 | } |
| 37 | |
| 38 | tensorflow::Status RunTPUBridge( |
| 39 | ModuleOp module, bool enable_logging, |
| 40 | llvm::function_ref<void(OpPassManager &pm)> pipeline_builder) { |
| 41 | PassManager bridge(module.getContext()); |
| 42 | |
| 43 | // Add logger to bridge passmanager. |
| 44 | if (enable_logging) |
| 45 | bridge.enableIRPrinting(std::make_unique<tensorflow::BridgeLoggerConfig>()); |
| 46 | |
| 47 | // Populate a passmanager with the list of passes that implement the bridge. |
| 48 | pipeline_builder(bridge); |
| 49 | |
| 50 | // Add set of passes to lower back to graph (from tf_executor). |
| 51 | AddGraphExportLoweringPasses(bridge); |
| 52 | |
| 53 | // Run the bridge on the module, in case of failure, the `diag_handler` |
| 54 | // converts MLIR errors emitted to the MLIRContext into a tensorflow::Status. |
| 55 | mlir::StatusScopedDiagnosticHandler diag_handler(module.getContext()); |
| 56 | LogicalResult result = bridge.run(module); |
| 57 | (void)result; |
| 58 | return diag_handler.ConsumeStatus(); |
| 59 | } |
| 60 | } // namespace |
| 61 | |
| 62 | void CreateTPUBridgePipeline(OpPassManager &pm) { |
no test coverage detected