Set TF_CUDA_COMPUTE_CAPABILITIES.
(environ_cp)
| 987 | |
| 988 | |
| 989 | def set_tf_cuda_compute_capabilities(environ_cp): |
| 990 | """Set TF_CUDA_COMPUTE_CAPABILITIES.""" |
| 991 | while True: |
| 992 | native_cuda_compute_capabilities = get_native_cuda_compute_capabilities( |
| 993 | environ_cp) |
| 994 | if not native_cuda_compute_capabilities: |
| 995 | default_cuda_compute_capabilities = _DEFAULT_CUDA_COMPUTE_CAPABILITIES |
| 996 | else: |
| 997 | default_cuda_compute_capabilities = native_cuda_compute_capabilities |
| 998 | |
| 999 | ask_cuda_compute_capabilities = ( |
| 1000 | 'Please specify a list of comma-separated ' |
| 1001 | 'CUDA compute capabilities you want to ' |
| 1002 | 'build with.\nYou can find the compute ' |
| 1003 | 'capability of your device at: ' |
| 1004 | 'https://developer.nvidia.com/cuda-gpus.\nPlease' |
| 1005 | ' note that each additional compute ' |
| 1006 | 'capability significantly increases your ' |
| 1007 | 'build time and binary size, and that ' |
| 1008 | 'TensorFlow only supports compute ' |
| 1009 | 'capabilities >= 3.5 [Default is: %s]: ' % |
| 1010 | default_cuda_compute_capabilities) |
| 1011 | tf_cuda_compute_capabilities = get_from_env_or_user_or_default( |
| 1012 | environ_cp, 'TF_CUDA_COMPUTE_CAPABILITIES', |
| 1013 | ask_cuda_compute_capabilities, default_cuda_compute_capabilities) |
| 1014 | # Check whether all capabilities from the input is valid |
| 1015 | all_valid = True |
| 1016 | # Remove all whitespace characters before splitting the string |
| 1017 | # that users may insert by accident, as this will result in error |
| 1018 | tf_cuda_compute_capabilities = ''.join(tf_cuda_compute_capabilities.split()) |
| 1019 | for compute_capability in tf_cuda_compute_capabilities.split(','): |
| 1020 | m = re.match('[0-9]+.[0-9]+', compute_capability) |
| 1021 | if not m: |
| 1022 | print('Invalid compute capability: %s' % compute_capability) |
| 1023 | all_valid = False |
| 1024 | else: |
| 1025 | ver = float(m.group(0)) |
| 1026 | if ver < 3.0: |
| 1027 | print('ERROR: TensorFlow only supports CUDA compute capabilities 3.0 ' |
| 1028 | 'and higher. Please re-specify the list of compute ' |
| 1029 | 'capabilities excluding version %s.' % ver) |
| 1030 | all_valid = False |
| 1031 | if ver < 3.5: |
| 1032 | print('WARNING: XLA does not support CUDA compute capabilities ' |
| 1033 | 'lower than 3.5. Disable XLA when running on older GPUs.') |
| 1034 | |
| 1035 | if all_valid: |
| 1036 | break |
| 1037 | |
| 1038 | # Reset and Retry |
| 1039 | environ_cp['TF_CUDA_COMPUTE_CAPABILITIES'] = '' |
| 1040 | |
| 1041 | # Set TF_CUDA_COMPUTE_CAPABILITIES |
| 1042 | environ_cp['TF_CUDA_COMPUTE_CAPABILITIES'] = tf_cuda_compute_capabilities |
| 1043 | write_action_env_to_bazelrc('TF_CUDA_COMPUTE_CAPABILITIES', |
| 1044 | tf_cuda_compute_capabilities) |
| 1045 | |
| 1046 |
no test coverage detected