From 68148b1bf34dc3365abebaacbc8210e6be380b75 Mon Sep 17 00:00:00 2001
From: "google-labs-jules[bot]"
 <161369871+google-labs-jules[bot]@users.noreply.github.com>
Date: Sun, 23 Nov 2025 08:19:07 +0000
Subject: [PATCH] Export TRITON_HIP_LLD_PATH in speedrun.sh for AMD ROCm

When running on AMD ROCm using `uv`-installed packages (`rocm-sdk-core`), the `ld.lld` linker is not in the default `/opt/rocm/llvm/bin/` location expected by `pytorch-triton-rocm`. This causes `InductorError` during `torch.compile`.

This change updates `speedrun.sh` to dynamically find the `ld.lld` binary within the active python environment's site-packages (`_rocm_sdk_core`) and export the `TRITON_HIP_LLD_PATH` environment variable, allowing Triton to locate the linker correctly.
---
 speedrun.sh | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/speedrun.sh b/speedrun.sh
index ebe625d..dd126da 100644
--- a/speedrun.sh
+++ b/speedrun.sh
@@ -49,6 +49,13 @@ if [ "$EXTRAS" == "amd" ]; then
     # Uninstalling triton may have deleted the shared 'triton' directory, breaking pytorch-triton-rocm.
     # Reinstall pytorch-triton-rocm to ensure it's intact.
     uv pip install --force-reinstall --index-url https://repo.amd.com/rocm/whl/gfx1151 pytorch-triton-rocm
+
+    # Find and export the path to ld.lld from rocm-sdk-core if available, as torch.compile/triton needs it
+    ROCM_LLD_PATH=$(python -c "import sysconfig; import os; p = f\"{sysconfig.get_paths()['purelib']}/_rocm_sdk_core/lib/llvm/bin/ld.lld\"; print(p) if os.path.exists(p) else print('')")
+    if [ -n "$ROCM_LLD_PATH" ]; then
+        export TRITON_HIP_LLD_PATH=$ROCM_LLD_PATH
+        echo "Exported TRITON_HIP_LLD_PATH=$TRITON_HIP_LLD_PATH"
+    fi
 fi
 
 # -----------------------------------------------------------------------------