| 53 | } |
| 54 | |
| 55 | void TF_EnableXLACompilation(TF_SessionOptions* options, unsigned char enable) { |
| 56 | tensorflow::ConfigProto& config = options->options.config; |
| 57 | auto* optimizer_options = |
| 58 | config.mutable_graph_options()->mutable_optimizer_options(); |
| 59 | if (enable) { |
| 60 | optimizer_options->set_global_jit_level(tensorflow::OptimizerOptions::ON_1); |
| 61 | |
| 62 | // These XLA flags are needed to trigger XLA properly from C (more generally |
| 63 | // non-Python) clients. If this API is called again with `enable` set to |
| 64 | // false, it is safe to keep these flag values as is. |
| 65 | tensorflow::MarkForCompilationPassFlags* flags = |
| 66 | tensorflow::GetMarkForCompilationPassFlags(); |
| 67 | flags->tf_xla_cpu_global_jit = true; |
| 68 | flags->tf_xla_min_cluster_size = 1; |
| 69 | } else { |
| 70 | optimizer_options->set_global_jit_level(tensorflow::OptimizerOptions::OFF); |
| 71 | } |
| 72 | } |
| 73 | |
| 74 | unsigned char TF_SetXlaEnableLazyCompilation(unsigned char enable) { |
| 75 | tensorflow::BuildXlaOpsPassFlags* flags = |