def get_cuda_compute_capability(source_from_url=False):
"""Retrieves CUDA compute capability based on the detected GPU type.
This function uses the `cuda_compute_capability` module to retrieve the
corresponding CUDA compute capability for the given GPU type.
Args:
source_from_url: Boolean deciding whether to source compute capability
from NVIDIA website or from a local golden file.
Returns:
List of all supported CUDA compute capabilities for the given GPU type.
e.g. ['3.5', '3.7']
"""
if not GPU_TYPE:
if FLAGS.debug:
print("Warning: GPU_TYPE is empty. "
"Make sure to call `get_gpu_type()` first.")
elif GPU_TYPE == "unknown":
if FLAGS.debug:
print("Warning: Unknown GPU is detected. "
"Skipping CUDA compute capability retrieval.")
else:
if source_from_url:
cuda_compute_capa = cuda_compute_capability.retrieve_from_web()
else:
cuda_compute_capa = cuda_compute_capability.retrieve_from_golden()
return cuda_compute_capa[GPU_TYPE]
return
def get_gpu_type():
"""Retrieves GPU type.
Returns:
String that is the name of the detected NVIDIA GPU.
e.g. 'Tesla K80'
'unknown' will be returned if detected GPU type is an unknown name.
Unknown name refers to any GPU name that is not specified in this page:
https://developer.nvidia.com/cuda-gpus
"""
global GPU_TYPE
key = "gpu_type_no_sudo"
gpu_dict = cuda_compute_capability.retrieve_from_golden()
out, err = run_shell_cmd(cmds_all[PLATFORM][key])
ret_val = out.split(b" ")
gpu_id = ret_val[0]
if err and FLAGS.debug:
print("Error in detecting GPU type:\n %s" % str(err))
if not isinstance(ret_val, list):
GPU_TYPE = "unknown"
return gpu_id, GPU_TYPE
else:
if "[" or "]" in ret_val[1]:
gpu_release = ret_val[1].replace(b"[", b"") + b" "
gpu_release += ret_val[2].replace(b"]", b"").strip(b"\n")
else:
gpu_release = six.ensure_str(ret_val[1]).replace("\n", " ")
if gpu_release not in gpu_dict:
GPU_TYPE = "unknown"
else:
GPU_TYPE = gpu_release
return gpu_id, GPU_TYPE