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

[Triton-MLIR][RUNTIME] Add /usr/bin/ptxas as a search path (#909)

Browse Source 
Make `ptxas` search a bit broader to include `/usr/bin/ptxas`, installed
by the lambda stack repo versions:
https://lambdalabs.com/lambda-stack-deep-learning-software
This commit is contained in:
Crutcher Dunnavant
2022-11-24 10:49:16 -08:00
committed by GitHub
parent ace7d28736
commit f98aed1258
1 changed files with 6 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
python/triton/compiler.py
View File

@@ -977,7 +977,12 @@ def ptx_get_version(cuda_version) -> int:
def path_to_ptxas():
prefixes = [os.environ.get("TRITON_PTXAS_PATH", ""), "", os.environ.get('CUDA_PATH', default_cuda_dir())]
prefixes = [
os.environ.get("TRITON_PTXAS_PATH", ""),
"",
"/usr",
os.environ.get('CUDA_PATH', default_cuda_dir())
]
for prefix in prefixes:
ptxas = os.path.join(prefix, "bin", "ptxas")
if os.path.exists(ptxas):