c74e5ffb65
ROCm packages are a runtime only dep for triton. triton-llvm always supports AMD GPU targets, so we can reduce how many different builds of triton are needed by teaching triton to better search for libamdhip64.so and ld.lld.
15 lines
582 B
Diff
15 lines
582 B
Diff
diff --git a/third_party/nvidia/backend/driver.py b/third_party/nvidia/backend/driver.py
|
|
index d088ec092..625de2db8 100644
|
|
--- a/third_party/nvidia/backend/driver.py
|
|
+++ b/third_party/nvidia/backend/driver.py
|
|
@@ -23,6 +23,9 @@ def libcuda_dirs():
|
|
if env_libcuda_path:
|
|
return [env_libcuda_path]
|
|
|
|
+ if os.path.exists("@libcudaStubsDir@"):
|
|
+ return ["@libcudaStubsDir@"]
|
|
+
|
|
libs = subprocess.check_output(["/sbin/ldconfig", "-p"]).decode()
|
|
# each line looks like the following:
|
|
# libcuda.so.1 (libc6,x86-64) => /lib/x86_64-linux-gnu/libcuda.so.1
|