From f98aed1258b0e52811e3f51fcba52c6b829beddd Mon Sep 17 00:00:00 2001 From: Crutcher Dunnavant Date: Thu, 24 Nov 2022 10:49:16 -0800 Subject: [PATCH] [Triton-MLIR][RUNTIME] Add /usr/bin/ptxas as a search path (#909) 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 --- python/triton/compiler.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/python/triton/compiler.py b/python/triton/compiler.py index 7eb30cb8e..42b47d3ad 100644 --- a/python/triton/compiler.py +++ b/python/triton/compiler.py @@ -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):