1#!/bin/bash 2set -x 3 4if [ -z "$1" ]; then 5 echo "Need wheel location argument" && exit 1 6fi 7 8WHEELHOUSE_DIR=$1 9PATCHELF_BIN=patchelf 10ROCM_LIB=backends/amd/lib 11ROCM_LD=backends/amd/llvm/bin 12PREFIX=triton 13fname_without_so_number() { 14 LINKNAME=$(echo $1 | sed -e 's/\.so.*/.so/g') 15 echo "$LINKNAME" 16} 17 18replace_needed_sofiles() { 19 find $1 -name '*.so*' -o -name 'ld.lld' | while read sofile; do 20 origname=$2 21 patchedname=$3 22 if [[ "$origname" != "$patchedname" ]]; then 23 set +e 24 origname=$($PATCHELF_BIN --print-needed $sofile | grep "$origname.*") 25 ERRCODE=$? 26 set -e 27 if [ "$ERRCODE" -eq "0" ]; then 28 echo "patching $sofile entry $origname to $patchedname" 29 $PATCHELF_BIN --replace-needed $origname $patchedname $sofile 30 fi 31 fi 32 done 33} 34 35mkdir -p "/tmp_dir" 36pushd /tmp_dir 37for pkg in /$WHEELHOUSE_DIR/*triton*.whl; do 38 echo "Modifying $pkg" 39 rm -rf tmp 40 mkdir -p tmp 41 cd tmp 42 cp $pkg . 43 unzip -q $(basename $pkg) 44 rm -f $(basename $pkg) 45 $PATCHELF_BIN --set-rpath ${LD_SO_RPATH:-'$ORIGIN:$ORIGIN/../../lib'} $PREFIX/$ROCM_LD/ld.lld 46 $PATCHELF_BIN --print-rpath $PREFIX/$ROCM_LD/ld.lld 47 # Modify libtriton.so as it sits in _C directory apart from its dependencies 48 find $PREFIX/_C -type f -name "*.so*" | while read sofile; do 49 echo "Setting rpath of $sofile" 50 $PATCHELF_BIN --set-rpath ${C_SO_RPATH:-'$ORIGIN:$ORIGIN/'../$ROCM_LIB} ${FORCE_RPATH:-} $sofile 51 $PATCHELF_BIN --print-rpath $sofile 52 done 53 54 # All included dependencies are included in a single lib directory 55 deps=() 56 deps_soname=() 57 while read sofile; do 58 echo "Setting rpath of $sofile to ${LIB_SO_RPATH:-'$ORIGIN'}" 59 $PATCHELF_BIN --set-rpath ${LIB_SO_RPATH:-'$ORIGIN'} ${FORCE_RPATH:-} $sofile 60 $PATCHELF_BIN --print-rpath $sofile 61 deps+=("$sofile") 62 deps_soname+=("$(basename $sofile)") 63 done < <(find $PREFIX/$ROCM_LIB -type f -name "*.so*") 64 65 patched=() 66 for filepath in "${deps[@]}"; do 67 filename=$(basename $filepath) 68 destpath=$PREFIX/$ROCM_LIB/$filename 69 if [[ "$filepath" != "$destpath" ]]; then 70 cp $filepath $destpath 71 fi 72 patchedpath=$(fname_without_so_number $destpath) 73 patchedname=$(basename $patchedpath) 74 if [[ "$destpath" != "$patchedpath" ]]; then 75 mv $destpath $patchedpath 76 fi 77 patched+=("$patchedname") 78 echo "Copied $filepath to $patchedpath" 79 done 80 81 # Go through all required shared objects and see if any of our other objects are dependants. If so, replace so.ver wth so 82 for ((i=0;i<${#deps[@]};++i)); do 83 echo "replacing "${deps_soname[i]} ${patched[i]} 84 replace_needed_sofiles $PREFIX/$ROCM_LIB ${deps_soname[i]} ${patched[i]} 85 replace_needed_sofiles $PREFIX/_C ${deps_soname[i]} ${patched[i]} 86 replace_needed_sofiles $PREFIX/$ROCM_LD ${deps_soname[i]} ${patched[i]} 87 done 88 89 # Re-bundle whl with so adjustments 90 zip -rqy $(basename $pkg) * 91 92 if [[ -z "${MANYLINUX_VERSION}" ]]; then 93 newpkg=$pkg 94 else 95 newpkg=$(echo $pkg | sed -e "s/\linux_x86_64/${MANYLINUX_VERSION}/g") 96 fi 97 98 # Remove original whl 99 rm -f $pkg 100 101 # Move rebuilt whl to original location with new name. 102 mv $(basename $pkg) $newpkg 103done 104