gaspersic/triton
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

[RUNTIME] Support environments with multiple cudalibs (#733)

Browse Source
This commit is contained in:
fdrocha
2022-10-03 19:36:24 +01:00
committed by GitHub
parent 4a2d3b7d79
commit 2b0f877fad
2 changed files with 9 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
python/test/unit/language/test_core.py
View File

@@ -666,7 +666,7 @@ def test_atomic_cas():
Lock = torch.zeros((1,), device='cuda', dtype=torch.int32) Lock = torch.zeros((1,), device='cuda', dtype=torch.int32)
change_value[(1,)](Lock) change_value[(1,)](Lock)
assert(Lock[0] == 1) assert (Lock[0] == 1)
# 2. only one block enters the critical section # 2. only one block enters the critical section
@triton.jit @triton.jit

14
python/triton/compiler.py
View File

@@ -1102,9 +1102,9 @@ class CacheManager:
@functools.lru_cache() @functools.lru_cache()
def libcuda_dir(): def libcuda_dirs():
loc = subprocess.check_output(["whereis", "libcuda.so"]).decode().strip().split()[-1] locs = subprocess.check_output(["whereis", "libcuda.so"]).decode().strip().split()[1:]
return os.path.dirname(loc) return [os.path.dirname(loc) for loc in locs]
@contextlib.contextmanager @contextlib.contextmanager
@@ -1118,7 +1118,7 @@ def quiet():
def _build(name, src, srcdir): def _build(name, src, srcdir):
cuda_lib_dir = libcuda_dir() cuda_lib_dirs = libcuda_dirs()
cu_include_dir = "/usr/local/cuda/include" cu_include_dir = "/usr/local/cuda/include"
suffix = sysconfig.get_config_var('EXT_SUFFIX') suffix = sysconfig.get_config_var('EXT_SUFFIX')
so = os.path.join(srcdir, '{name}{suffix}'.format(name=name, suffix=suffix)) so = os.path.join(srcdir, '{name}{suffix}'.format(name=name, suffix=suffix))
@@ -1130,12 +1130,14 @@ def _build(name, src, srcdir):
gcc = shutil.which("gcc") gcc = shutil.which("gcc")
cc = gcc if gcc is not None else clang cc = gcc if gcc is not None else clang
py_include_dir = get_paths()["include"] py_include_dir = get_paths()["include"]
ret = subprocess.check_call([cc, src, "-O3", f"-I{cu_include_dir}", f"-I{py_include_dir}", f"-I{srcdir}", "-shared", "-fPIC", f"-L{cuda_lib_dir}", "-lcuda", "-o", so]) cc_cmd = [cc, src, "-O3", f"-I{cu_include_dir}", f"-I{py_include_dir}", f"-I{srcdir}", "-shared", "-fPIC", "-lcuda", "-o", so]
cc_cmd += [f"-L{dir}" for dir in cuda_lib_dirs]
ret = subprocess.check_call(cc_cmd)
if ret == 0: if ret == 0:
return so return so
# fallback on setuptools # fallback on setuptools
extra_compile_args = [] extra_compile_args = []
library_dirs = [cuda_lib_dir] library_dirs = cuda_lib_dirs
include_dirs = [srcdir, cu_include_dir] include_dirs = [srcdir, cu_include_dir]
libraries = ['cuda'] libraries = ['cuda']
# extra arguments # extra arguments