OSSForks/pytorch
1
0
Fork 0
mirror of https://github.com/zebrajr/pytorch.git synced 2025-12-06 12:20:52 +01:00
Code Issues Actions 38 Packages Projects Releases Wiki Activity

Make triton build ROCm library version-agnostic (#158408)

Browse Source 
Fixes maintenance of triton packaging script when library versions change from one ROCm version to next.

Pull Request resolved: https://github.com/pytorch/pytorch/pull/158408
Approved by: https://github.com/jeffdaily

Co-authored-by: Ethan Wee <Ethan.Wee@amd.com>
This commit is contained in:
Jithun Nair 2025-08-13 19:49:23 +00:00 committed by PyTorch MergeBot
parent 70ccdec44b
commit 4cde0acc0e
2 changed files with 14 additions and 33 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

31
.github/scripts/amd/package_triton_wheel.sh vendored
View File

@ -1,3 +1,4 @@
#!/bin/bash
set -ex set -ex
# Set ROCM_HOME isn't available, use ROCM_PATH if set or /opt/rocm # Set ROCM_HOME isn't available, use ROCM_PATH if set or /opt/rocm
@ -50,29 +51,15 @@ do
cp $lib $TRITON_ROCM_DIR/lib/ cp $lib $TRITON_ROCM_DIR/lib/
done done
# Required ROCm libraries
if [[ "${MAJOR_VERSION}" == "6" ]]; then
libamdhip="libamdhip64.so.6"
else
libamdhip="libamdhip64.so.5"
fi
# Required ROCm libraries - ROCm 6.0 # Required ROCm libraries - ROCm 6.0
ROCM_SO=( ROCM_SO=(
"${libamdhip}" "libamdhip64.so"
"libhsa-runtime64.so.1" "libhsa-runtime64.so"
"libdrm.so.2" "libdrm.so"
"libdrm_amdgpu.so.1" "libdrm_amdgpu.so"
"libamd_comgr.so"
"librocprofiler-register.so"
) )
if [[ $ROCM_INT -ge 60400 ]]; then
ROCM_SO+=("libamd_comgr.so.3")
else
ROCM_SO+=("libamd_comgr.so.2")
fi
if [[ $ROCM_INT -ge 60100 ]]; then
ROCM_SO+=("librocprofiler-register.so.0")
fi
for lib in "${ROCM_SO[@]}" for lib in "${ROCM_SO[@]}"
do do
@ -94,10 +81,6 @@ do
fi fi
cp $file_path $TRITON_ROCM_DIR/lib cp $file_path $TRITON_ROCM_DIR/lib
# When running locally, and not building a wheel, we need to satisfy shared objects requests that don't look for versions
LINKNAME=$(echo $lib | sed -e 's/\.so.*/.so/g')
ln -sf $lib $TRITON_ROCM_DIR/lib/$LINKNAME
done done
# Copy Include Files # Copy Include Files

16
.github/scripts/amd/patch_triton_wheel.sh vendored
View File

@ -19,15 +19,13 @@ replace_needed_sofiles() {
find $1 -name '*.so*' -o -name 'ld.lld' | while read sofile; do find $1 -name '*.so*' -o -name 'ld.lld' | while read sofile; do
origname=$2 origname=$2
patchedname=$3 patchedname=$3
if [[ "$origname" != "$patchedname" ]]; then set +e
set +e origname=$($PATCHELF_BIN --print-needed $sofile | grep "$origname.*")
origname=$($PATCHELF_BIN --print-needed $sofile | grep "$origname.*") ERRCODE=$?
ERRCODE=$? set -e
set -e if [ "$ERRCODE" -eq "0" ]; then
if [ "$ERRCODE" -eq "0" ]; then echo "patching $sofile entry $origname to $patchedname"
echo "patching $sofile entry $origname to $patchedname" $PATCHELF_BIN --replace-needed $origname $patchedname $sofile
$PATCHELF_BIN --replace-needed $origname $patchedname $sofile
fi
fi fi
done done
} }