From 85408f52667c75b74ef2bdbcccb9d080eea3d2af Mon Sep 17 00:00:00 2001
From: Stephan Seitz <stephan.seitz@fau.de>
Date: Fri, 20 Sep 2019 22:33:07 +0200
Subject: [PATCH] Enable setting NVCC_BINARY by environment variable
---
src/pystencils_autodiff/tensorflow_jit.py | 8 ++++++--
1 file changed, 6 insertions(+), 2 deletions(-)
diff --git a/src/pystencils_autodiff/tensorflow_jit.py b/src/pystencils_autodiff/tensorflow_jit.py
index 909b980..3b90107 100644
--- a/src/pystencils_autodiff/tensorflow_jit.py
+++ b/src/pystencils_autodiff/tensorflow_jit.py
@@ -14,13 +14,17 @@ import sysconfig
from os.path import exists, join
import p_tqdm
-
import pystencils
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
_hash = hashlib.md5
+if 'NVCC_BINARY' in os.environ:
+ NVCC_BINARY = os.environ['NVCC_BINARY']
+else:
+ NVCC_BINARY = 'nvcc'
# TODO: msvc
if get_compiler_config()['os'] != 'windows':
@@ -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']
if use_nvcc:
- command_prefix = [nvcc,
+ command_prefix = [NVCC_BINARY,
'--expt-relaxed-constexpr',
'-ccbin',
get_compiler_config()['tensorflow_host_compiler'],
--
GitLab