| 100 | } |
| 101 | |
| 102 | TF_Buffer* TF_CreateConfig(unsigned char enable_xla_compilation, |
| 103 | unsigned char gpu_memory_allow_growth, |
| 104 | unsigned int num_cpu_devices) { |
| 105 | tensorflow::ConfigProto config; |
| 106 | auto* optimizer_options = |
| 107 | config.mutable_graph_options()->mutable_optimizer_options(); |
| 108 | if (enable_xla_compilation) { |
| 109 | optimizer_options->set_global_jit_level(tensorflow::OptimizerOptions::ON_1); |
| 110 | |
| 111 | // These XLA flags are needed to trigger XLA properly from C (more generally |
| 112 | // non-Python) clients. If this API is called again with `enable` set to |
| 113 | // false, it is safe to keep these flag values as is. |
| 114 | tensorflow::MarkForCompilationPassFlags* flags = |
| 115 | tensorflow::GetMarkForCompilationPassFlags(); |
| 116 | flags->tf_xla_cpu_global_jit = true; |
| 117 | flags->tf_xla_min_cluster_size = 1; |
| 118 | } else { |
| 119 | optimizer_options->set_global_jit_level(tensorflow::OptimizerOptions::OFF); |
| 120 | } |
| 121 | |
| 122 | auto* gpu_options = config.mutable_gpu_options(); |
| 123 | gpu_options->set_allow_growth(gpu_memory_allow_growth); |
| 124 | |
| 125 | (*config.mutable_device_count())["CPU"] = num_cpu_devices; |
| 126 | |
| 127 | // TODO(b/113217601): This is needed for EagerContext::runner_ to use a |
| 128 | // threadpool, so that we avoid the possibility of running the runner_ in the |
| 129 | // threadpool of GPU event mgr, as that can trigger more callbacks to be |
| 130 | // scheduled on that same threadpool, causing a deadlock in cases where the |
| 131 | // caller of event_mgr->ThenExecute() blocks on the completion of the callback |
| 132 | // (as in the case of ConstOp kernel creation on GPU, which involves copying a |
| 133 | // CPU tensor to GPU). |
| 134 | // Setting a larger thread pool does not help with the Swift caller, as we use |
| 135 | // a different TFE context for each thread of execution (for running graph |
| 136 | // functions, and their send/recvs corountines). |
| 137 | config.set_inter_op_parallelism_threads(1); |
| 138 | |
| 139 | TF_Buffer* ret = TF_NewBuffer(); |
| 140 | TF_CHECK_OK(MessageToBuffer(config, ret)); |
| 141 | return ret; |
| 142 | } |
| 143 | |
| 144 | TF_Buffer* TF_CreateRunOptions(unsigned char enable_full_trace) { |
| 145 | tensorflow::RunOptions options; |
no test coverage detected