MCPcopy Create free account
hub / github.com/DeepRec-AI/DeepRec / TF_CreateConfig

Function TF_CreateConfig

tensorflow/c/c_api_experimental.cc:102–142  ·  view source on GitHub ↗

Source from the content-addressed store, hash-verified

100}
101
102TF_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
144TF_Buffer* TF_CreateRunOptions(unsigned char enable_full_trace) {
145 tensorflow::RunOptions options;

Callers 1

Calls 3

TF_NewBufferFunction · 0.85
MessageToBufferFunction · 0.85

Tested by

no test coverage detected