Skip to content
Snippets Groups Projects
Commit 85408f52 authored by Stephan Seitz's avatar Stephan Seitz
Browse files

Enable setting NVCC_BINARY by environment variable

parent ff61f430
No related branches found
No related tags found
No related merge requests found
Pipeline #18205 failed
...@@ -14,13 +14,17 @@ import sysconfig ...@@ -14,13 +14,17 @@ import sysconfig
from os.path import exists, join from os.path import exists, join
import p_tqdm import p_tqdm
import pystencils import pystencils
from pystencils.cpu.cpujit import get_cache_config, get_compiler_config, get_pystencils_include_path from pystencils.cpu.cpujit import get_cache_config, get_compiler_config, get_pystencils_include_path
from pystencils_autodiff._file_io import read_file, write_file from pystencils_autodiff._file_io import read_file, write_file
_hash = hashlib.md5 _hash = hashlib.md5
if 'NVCC_BINARY' in os.environ:
NVCC_BINARY = os.environ['NVCC_BINARY']
else:
NVCC_BINARY = 'nvcc'
# TODO: msvc # TODO: msvc
if get_compiler_config()['os'] != 'windows': if get_compiler_config()['os'] != 'windows':
...@@ -120,7 +124,7 @@ def compile_file(file, use_nvcc=False, nvcc='nvcc', overwrite_destination_file=T ...@@ -120,7 +124,7 @@ def compile_file(file, use_nvcc=False, nvcc='nvcc', overwrite_destination_file=T
get_compiler_config()['tensorflow_host_compiler'] = get_compiler_config()['command'] get_compiler_config()['tensorflow_host_compiler'] = get_compiler_config()['command']
if use_nvcc: if use_nvcc:
command_prefix = [nvcc, command_prefix = [NVCC_BINARY,
'--expt-relaxed-constexpr', '--expt-relaxed-constexpr',
'-ccbin', '-ccbin',
get_compiler_config()['tensorflow_host_compiler'], get_compiler_config()['tensorflow_host_compiler'],
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Please register or to comment