| 27 | namespace { |
| 28 | static absl::once_flag flag_init; |
| 29 | static void GetCudaContextsCount(int ordinal, int64* cuda_contexts_count) { |
| 30 | *cuda_contexts_count = 1; |
| 31 | #ifdef GOOGLE_CUDA |
| 32 | gpu::GpuDriver::Init(); |
| 33 | gpu::GpuDeviceHandle device; |
| 34 | auto status = gpu::GpuDriver::GetDevice(ordinal, &device); |
| 35 | if (status.ok()) { |
| 36 | int cc_major = 0, cc_minor = 0; |
| 37 | gpu::GpuDriver::GetComputeCapability(&cc_major, &cc_minor, device); |
| 38 | if (cc_major >= 7) { |
| 39 | int64 ctx_count = 4; |
| 40 | tensorflow::ReadInt64FromEnvVar("CONTEXTS_COUNT_PER_GPU", 4, &ctx_count); |
| 41 | if (ctx_count > 0) { |
| 42 | *cuda_contexts_count = ctx_count; |
| 43 | } |
| 44 | LOG(INFO) << "User set " << *cuda_contexts_count << " cuda context for each gpu."; |
| 45 | } else { |
| 46 | LOG(ERROR) << "Enable MPS required cc_major >= 7, now is " << cc_major; |
| 47 | } |
| 48 | } else { |
| 49 | LOG(ERROR) << "Cuda get device " << ordinal << " error: " |
| 50 | << status.error_message(); |
| 51 | } |
| 52 | #endif // GOOGLE_CUDA |
| 53 | } |
| 54 | } // end namespace |
| 55 | |
| 56 | port::StatusOr<StreamExecutor*> ExecutorCache::GetOrCreate( |
nothing calls this directly
no test coverage detected